Adding initial RETCALL to all section and paragraph endings in IR

avishek-sen-gupta · Nov 19, 2024 · e84c1fe · e84c1fe
1 parent 70cfe8b
commit e84c1fe
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 13 deletions.
diff --git a/smojol-core/src/main/java/org/smojol/common/navigation/FlowNodeNavigator.java b/smojol-core/src/main/java/org/smojol/common/navigation/FlowNodeNavigator.java
@@ -22,6 +22,10 @@ public <T extends FlowNode> T findByType(Class<T> type) {
         return (T) findByCondition(fn -> fn.getClass() == type);
     }
 
+    public <T extends FlowNode> List<T> findAllByType(Class<T> type) {
+        return findAllByCondition(fn -> fn.getClass() == type).stream().map(n -> (T) n).toList();
+    }
+
     private FlowNode searchRecursively(FlowNode current, FlowNodeCondition c) {
         if (c.apply(current)) return current;
         for (FlowNode child : current.astChildren()) {

diff --git a/smojol-core/src/main/java/org/smojol/common/transpiler/RetCallTranspilerNode.java b/smojol-core/src/main/java/org/smojol/common/transpiler/RetCallTranspilerNode.java
@@ -6,7 +6,7 @@
 public class RetCallTranspilerNode extends TranspilerNode {
     private final String fallthroughTarget;
 
-    public RetCallTranspilerNode(String fallthroughTarget, TranspilerCodeBlockNode body) {
+    public RetCallTranspilerNode(String fallthroughTarget) {
         super(ImmutableList.of(), ImmutableList.of(SemanticCategory.BLOCK_BOUNDARY));
         this.fallthroughTarget = fallthroughTarget;
     }

diff --git a/...oolkit/src/main/java/org/smojol/toolkit/analysis/task/transpiler/SectionParagraphMap.java b/...oolkit/src/main/java/org/smojol/toolkit/analysis/task/transpiler/SectionParagraphMap.java
@@ -6,25 +6,44 @@
 import org.smojol.common.ast.FlowNode;
 import org.smojol.common.ast.FlowNodeType;
 import org.smojol.common.navigation.FlowNodeNavigator;
+import org.smojol.toolkit.ast.ParagraphFlowNode;
+import org.smojol.toolkit.ast.SectionFlowNode;
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 
 @Getter
 public class SectionParagraphMap {
-    private final Map<FlowNode, List<FlowNode>> sectionParagraphsPairs;
-    private final Map<FlowNode, FlowNode> paragraphToSectionMap;
+    private final Map<SectionFlowNode, List<ParagraphFlowNode>> sectionParagraphsMap;
+    private final Map<ParagraphFlowNode, SectionFlowNode> paragraphToSectionMap;
+    private final List<SectionFlowNode> sectionFlowNodes;
 
     public SectionParagraphMap(FlowNode flowRoot) {
-        List<FlowNode> sectionFlowNodes = new FlowNodeNavigator(flowRoot).findAllByCondition(fn -> fn.type() == FlowNodeType.SECTION);
-        sectionParagraphsPairs = sectionFlowNodes.stream()
-                .map(sfn -> (Pair<FlowNode, List<FlowNode>>) ImmutablePair.of(sfn, new FlowNodeNavigator(sfn).findAllByCondition(n -> n.type() == FlowNodeType.PARAGRAPH)))
+        sectionFlowNodes = new FlowNodeNavigator(flowRoot).findAllByType(SectionFlowNode.class);
+        sectionParagraphsMap = sectionFlowNodes.stream()
+                .map(sfn -> (Pair<SectionFlowNode, List<ParagraphFlowNode>>) ImmutablePair.of(sfn, new FlowNodeNavigator(sfn).findAllByType(ParagraphFlowNode.class)))
                 .collect(Collectors.toMap(Pair::getLeft, Pair::getRight));
         paragraphToSectionMap = sectionFlowNodes.stream().flatMap(sfn ->
                 new FlowNodeNavigator(sfn).findAllByCondition(n -> n.type() == FlowNodeType.PARAGRAPH).stream()
-                        .map(para -> ImmutablePair.of(para, sfn))
+                        .map(para -> ImmutablePair.of((ParagraphFlowNode) para, sfn))
         ).collect(Collectors.toMap(Pair::getLeft, Pair::getRight));
 
     }
+
+    public Optional<SectionFlowNode> nextSection(SectionFlowNode n) {
+        Optional<SectionFlowNode> flowNode = sectionFlowNodes.stream().filter(sn -> sn == n).findFirst();
+        if (flowNode.isEmpty()) return flowNode;
+        int nodeIndex = sectionFlowNodes.indexOf(flowNode.get());
+        return nodeIndex == sectionFlowNodes.size() - 1 ? Optional.empty() : Optional.of(sectionFlowNodes.get(nodeIndex + 1));
+    }
+
+    public Optional<ParagraphFlowNode> nextParagraph(ParagraphFlowNode n) {
+        SectionFlowNode sectionFlowNode = paragraphToSectionMap.get(n);
+        List<ParagraphFlowNode> parasForSection = sectionParagraphsMap.get(sectionFlowNode);
+        int paraIndex = parasForSection.indexOf(n);
+        if (paraIndex == -1) return Optional.empty();
+        return paraIndex == parasForSection.size() - 1 ? Optional.empty() : Optional.of(parasForSection.get(paraIndex + 1));
+    }
 }
diff --git a/smojol-toolkit/src/main/java/org/smojol/toolkit/ast/ParagraphFlowNode.java b/smojol-toolkit/src/main/java/org/smojol/toolkit/ast/ParagraphFlowNode.java
@@ -1,6 +1,7 @@
 package org.smojol.toolkit.ast;
 
 import org.antlr.v4.runtime.tree.ParseTree;
+import org.eclipse.lsp.cobol.core.CobolParser;
 import org.smojol.common.ast.FlowNode;
 import org.smojol.common.ast.FlowNodeService;
 import org.smojol.common.ast.FlowNodeType;
@@ -41,4 +42,9 @@ protected CobolVmSignal continueOrAbort(CobolVmSignal defaultSignal, CobolInterp
         if (defaultSignal == CobolVmSignal.TERMINATE || defaultSignal == CobolVmSignal.EXIT_PERFORM) return defaultSignal;
         return next(defaultSignal, interpreter, nodeService);
     }
+
+    @Override
+    public String label() {
+        return ((CobolParser.ParagraphContext) executionContext).paragraphDefinitionName().getText();
+    }
 }
diff --git a/smojol-toolkit/src/main/java/org/smojol/toolkit/ast/SectionFlowNode.java b/smojol-toolkit/src/main/java/org/smojol/toolkit/ast/SectionFlowNode.java
@@ -1,6 +1,7 @@
 package org.smojol.toolkit.ast;
 
 import org.antlr.v4.runtime.tree.ParseTree;
+import org.eclipse.lsp.cobol.core.CobolParser;
 import org.smojol.common.ast.FlowNode;
 import org.smojol.common.ast.FlowNodeService;
 import org.smojol.common.ast.FlowNodeType;
@@ -31,4 +32,9 @@ protected CobolVmSignal continueOrAbort(CobolVmSignal defaultSignal, CobolInterp
         if (defaultSignal == CobolVmSignal.TERMINATE || defaultSignal == CobolVmSignal.EXIT_PERFORM) return defaultSignal;
         return next(defaultSignal, interpreter, nodeService);
     }
+
+    @Override
+    public String label() {
+        return ((CobolParser.ProcedureSectionContext) executionContext).procedureSectionHeader().sectionName().getText();
+    }
 }
diff --git a/...t/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java b/...t/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java
@@ -3,7 +3,10 @@
 import com.google.common.collect.ImmutableMap;
 import org.smojol.common.ast.FlowNode;
 import org.smojol.common.ast.FlowNodeType;
+import org.smojol.common.ast.NullFlowNode;
 import org.smojol.common.transpiler.LabelledTranspilerCodeBlockNode;
+import org.smojol.common.transpiler.NullTranspilerNode;
+import org.smojol.common.transpiler.RetCallTranspilerNode;
 import org.smojol.common.transpiler.TranspilerNode;
 import org.smojol.common.vm.structure.CobolDataStructure;
 import org.smojol.toolkit.analysis.task.transpiler.SectionParagraphMap;
@@ -13,23 +16,29 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Stream;
 
 public class LabelledTranspilerCodeBlockNodeBuilder {
     public static TranspilerNode build(SectionFlowNode n, CobolDataStructure dataStructures, SectionParagraphMap sectionParagraphMap) {
-        return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.SECTION), sectionParagraphMap);
+        FlowNode next = sectionParagraphMap.nextSection(n).get();
+        return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.SECTION), sectionParagraphMap, next);
     }
 
-    private static LabelledTranspilerCodeBlockNode labelledBlock(FlowNode n, CobolDataStructure dataStructures, Map<String, Object> properties, SectionParagraphMap sectionParagraphMap) {
-        List<TranspilerNode> childTranspilerNodes = n.astChildren().stream().map(child -> TranspilerTreeBuilder.flowToTranspiler(child, dataStructures, sectionParagraphMap)).toList();
-        return new LabelledTranspilerCodeBlockNode(n.name(), childTranspilerNodes, properties);
+    private static LabelledTranspilerCodeBlockNode labelledBlock(FlowNode n, CobolDataStructure dataStructures, Map<String, Object> properties, SectionParagraphMap sectionParagraphMap, FlowNode next) {
+        Stream<TranspilerNode> childTranspilerNodes = n.astChildren().stream().map(child -> TranspilerTreeBuilder.flowToTranspiler(child, dataStructures, sectionParagraphMap));
+        List<TranspilerNode> childNodesWithRetCall = Stream.concat(childTranspilerNodes, next instanceof NullTranspilerNode
+                ? Stream.of()
+                : Stream.of(new RetCallTranspilerNode(next.label()))).toList();
+        return new LabelledTranspilerCodeBlockNode(n.label(), childNodesWithRetCall, properties);
     }
 
     public static TranspilerNode build(ParagraphFlowNode n, CobolDataStructure dataStructures, SectionParagraphMap sectionParagraphMap) {
-        return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.PARAGRAPH), sectionParagraphMap);
+        ParagraphFlowNode next = sectionParagraphMap.nextParagraph(n).get();
+        return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.PARAGRAPH), sectionParagraphMap, next);
     }
 
     public static TranspilerNode build(ProcedureDivisionBodyFlowNode n, CobolDataStructure dataStructures, SectionParagraphMap sectionParagraphMap) {
-        return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.PROCEDURE_DIVISION_BODY), sectionParagraphMap);
+        return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.PROCEDURE_DIVISION_BODY), sectionParagraphMap, new NullFlowNode());
     }
 
 //    public static TranspilerNode build(ParagraphsFlowNode n, CobolDataStructure dataStructures) {