fix: fix folding when IfNode has a copynode child

Signed-off-by: Aman Prashant <aman.prashant@broadcom.com>
ap891843 · Feb 12, 2024 · 0c9717d · 0c9717d
1 parent d1e902f
commit 0c9717d
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 18 deletions.
diff --git a/server/engine/src/main/java/org/eclipse/lsp/cobol/service/DocumentServiceHelper.java b/server/engine/src/main/java/org/eclipse/lsp/cobol/service/DocumentServiceHelper.java
@@ -19,6 +19,7 @@
 import static org.eclipse.lsp.cobol.common.model.tree.Node.hasType;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.IntSummaryStatistics;
@@ -74,7 +75,7 @@ public static Set<FoldingRange> getFoldingRange(@NonNull Node rootNode, @NonNull
             node ->
                 node.getLocality().getRange().getStart().getLine()
                     != node.getLocality().getRange().getEnd().getLine())
-        .map(DocumentServiceHelper::getFoldingRange)
+        .map(node -> DocumentServiceHelper.getFoldingRanges(node, uri))
         .flatMap(List::stream)
         .collect(Collectors.toSet());
   }
@@ -91,34 +92,38 @@ private static FoldingRange getFoldingRange(List<Node> node) {
     return new FoldingRange(start.getMin(), end.getMax());
   }
 
-  private static List<FoldingRange> getFoldingRange(IfNode node) {
+  private static List<FoldingRange> getFoldingRange(IfNode node, String uri) {
     List<Node> ifThenStatements =
-        node.getChildren().stream().filter(n -> !(n instanceof IfElseNode)).collect(toList());
-    return ImmutableList.of(getFoldingRange(ifThenStatements));
+        node.getChildren().stream()
+            .filter(n -> !(n instanceof IfElseNode))
+            .filter(nod -> nod.getLocality().getUri().equals(uri))
+            .collect(toList());
+    return Lists.newArrayList(getFoldingRange(ifThenStatements));
   }
 
-  private static List<FoldingRange> getFoldingRange(EvaluateNode node) {
+  private static List<FoldingRange> getFoldingRange(EvaluateNode node, String uri) {
     Map<Node, List<Node>> accumulator = new HashMap<>();
     Node lastEvaluateNode = null;
     for (Node child : node.getChildren()) {
       if (child instanceof EvaluateWhenNode || child instanceof EvaluateWhenOtherNode) {
-        lastEvaluateNode =  child;
+        lastEvaluateNode = child;
         accumulator.putIfAbsent(lastEvaluateNode, new ArrayList<>());
       }
       if (lastEvaluateNode == null) continue;
-      accumulator.get(lastEvaluateNode).add(child);
+      if (child.getLocality().getUri().equals(uri)) accumulator.get(lastEvaluateNode).add(child);
     }
     return accumulator.values().stream()
-        .map(DocumentServiceHelper::getFoldingRange)
-        .collect(toList());
+            .map(DocumentServiceHelper::getFoldingRange)
+            .collect(toList());
   }
 
-  private static List<FoldingRange> getFoldingRange(Node node) {
+  private static List<FoldingRange> getFoldingRanges(Node node, String uri) {
     if (node instanceof IfNode) {
-      return getFoldingRange((IfNode) node);
+      return getFoldingRange((IfNode) node, uri);
     }
     if (node instanceof EvaluateNode) {
-      List<FoldingRange> foldingRange = getFoldingRange((EvaluateNode) node);
+      List<FoldingRange> foldingRange = getFoldingRange((EvaluateNode) node, uri);
+
       foldingRange.add(getFoldingRangeFromNode(node));
       return foldingRange;
     }

diff --git a/server/engine/src/test/java/org/eclipse/lsp/cobol/service/DocumentServiceHelperTest.java b/server/engine/src/test/java/org/eclipse/lsp/cobol/service/DocumentServiceHelperTest.java
@@ -18,14 +18,10 @@
 
 import java.util.Set;
 import org.eclipse.lsp.cobol.common.model.Locality;
-import org.eclipse.lsp.cobol.common.model.tree.EvaluateNode;
-import org.eclipse.lsp.cobol.common.model.tree.EvaluateWhenNode;
-import org.eclipse.lsp.cobol.common.model.tree.ExitNode;
-import org.eclipse.lsp.cobol.common.model.tree.IfNode;
-import org.eclipse.lsp.cobol.common.model.tree.Node;
-import org.eclipse.lsp.cobol.common.model.tree.RootNode;
+import org.eclipse.lsp.cobol.common.model.tree.*;
 import org.eclipse.lsp.cobol.common.model.tree.variable.QualifiedReferenceNode;
 import org.eclipse.lsp4j.FoldingRange;
+import org.eclipse.lsp4j.Location;
 import org.eclipse.lsp4j.Position;
 import org.eclipse.lsp4j.Range;
 import org.junit.jupiter.api.Disabled;
@@ -91,12 +87,31 @@ void getFoldingRange_whenEvaluateStatementIsPresent() {
     ifNode.addChild(new QualifiedReferenceNode(Locality.builder()
             .range(new Range(new Position(3, 0), new Position(4, 10)))
             .uri(DOCUMENT_URI).build()));
+    CopyNode copyNode = new CopyNode(
+            Locality.builder()
+                    .range(new Range(new Position(5, 0), new Position(5, 10)))
+                    .uri(DOCUMENT_URI)
+                    .build(),
+            new Location(),
+            "",
+            "");
+    PerformNode performNode = new PerformNode(Locality.builder()
+            .range(new Range(new Position(0, 0), new Position(12, 10)))
+            .uri("file:///c:/workspace/copy.cpy")
+            .build());
+    copyNode.addChild(performNode);
+    ifNode.addChild(performNode);
+    ifNode.addChild(copyNode);
     evaluateWhenNode.addChild(ifNode);
     evaluateNode.addChild(evaluateWhenNode);
     evaluateNode.addChild(evaluateWhenNode2);
     rootNode.addChild(evaluateNode);
     Set<FoldingRange> foldingRange = DocumentServiceHelper.getFoldingRange(rootNode, DOCUMENT_URI);
     assertEquals(4, foldingRange.size());
+    assertTrue(foldingRange.contains(new FoldingRange(5, 6))); // evaluateWhenNode2
+    assertTrue(foldingRange.contains(new FoldingRange(0, 10))); // evaluateNode
+    assertTrue(foldingRange.contains(new FoldingRange(3, 5))); // if node
+    assertTrue(foldingRange.contains(new FoldingRange(2, 4))); // evaluateWhenNode
   }
 
   private static Node getRootNode() {