Fixing unit test

avishek-sen-gupta · Nov 20, 2024 · 279c41c · 279c41c
1 parent 720159f
commit 279c41c
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/...t/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java b/...t/src/main/java/org/smojol/toolkit/transpiler/LabelledTranspilerCodeBlockNodeBuilder.java
@@ -5,7 +5,6 @@
 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;
@@ -16,24 +15,27 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Stream;
 
 public class LabelledTranspilerCodeBlockNodeBuilder {
     public static TranspilerNode build(SectionFlowNode n, CobolDataStructure dataStructures, SectionParagraphMap sectionParagraphMap) {
-        FlowNode next = sectionParagraphMap.nextSection(n).get();
+        Optional<SectionFlowNode> maybeNext = sectionParagraphMap.nextSection(n);
+        FlowNode next = maybeNext.isPresent() ? maybeNext.get() : new NullFlowNode();
         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, FlowNode next) {
         Stream<TranspilerNode> childTranspilerNodes = n.astChildren().stream().map(child -> TranspilerTreeBuilder.flowToTranspiler(child, dataStructures, sectionParagraphMap));
-        List<TranspilerNode> childNodesWithRetCall = Stream.concat(childTranspilerNodes, next instanceof NullTranspilerNode
+        List<TranspilerNode> childNodesWithRetCall = Stream.concat(childTranspilerNodes, next instanceof NullFlowNode
                 ? 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) {
-        ParagraphFlowNode next = sectionParagraphMap.nextParagraph(n).get();
+        Optional<ParagraphFlowNode> maybeNext = sectionParagraphMap.nextParagraph(n);
+        FlowNode next = maybeNext.isPresent() ? maybeNext.get() : new NullFlowNode();
         return labelledBlock(n, dataStructures, ImmutableMap.of("type", FlowNodeType.PARAGRAPH), sectionParagraphMap, next);
     }