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 index 43b31493..a33d19dc 100644 --- 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 findByType(Class type) { return (T) findByCondition(fn -> fn.getClass() == type); } + public List findAllByType(Class 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 index fc32ce9f..421dcecd 100644 --- 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/smojol-toolkit/src/main/java/org/smojol/toolkit/analysis/task/transpiler/SectionParagraphMap.java b/smojol-toolkit/src/main/java/org/smojol/toolkit/analysis/task/transpiler/SectionParagraphMap.java index 6783fbcc..ca47c5b0 100644 --- a/smojol-toolkit/src/main/java/org/smojol/toolkit/analysis/task/transpiler/SectionParagraphMap.java +++ b/smojol-toolkit/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> sectionParagraphsPairs; - private final Map paragraphToSectionMap; + private final Map> sectionParagraphsMap; + private final Map paragraphToSectionMap; + private final List sectionFlowNodes; public SectionParagraphMap(FlowNode flowRoot) { - List sectionFlowNodes = new FlowNodeNavigator(flowRoot).findAllByCondition(fn -> fn.type() == FlowNodeType.SECTION); - sectionParagraphsPairs = sectionFlowNodes.stream() - .map(sfn -> (Pair>) 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>) 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 nextSection(SectionFlowNode n) { + Optional 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 nextParagraph(ParagraphFlowNode n) { + SectionFlowNode sectionFlowNode = paragraphToSectionMap.get(n); + List 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 index bb1ab1c3..9c7f2621 100644 --- 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 index 262e5322..6aa7d08b 100644 --- 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/smojol-toolkit/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java b/smojol-toolkit/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java index 72823f10..91f1ecaa 100644 --- a/smojol-toolkit/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java +++ b/smojol-toolkit/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 properties, SectionParagraphMap sectionParagraphMap) { - List 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 properties, SectionParagraphMap sectionParagraphMap, FlowNode next) { + Stream childTranspilerNodes = n.astChildren().stream().map(child -> TranspilerTreeBuilder.flowToTranspiler(child, dataStructures, sectionParagraphMap)); + List 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) {