Explicitly start basic block instead of setting label

1 files changed, 26 insertions, 25 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
index 6a90e4b1..576ac771 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -424,15 +424,13 @@ public class LlvmPatioCodeGen {
                 var loopEndLabel = loopEndLabelResult.value();
                 context = loopEndLabelResult.context();
 
-                context = context.addTerminatorInstruction(new LlvmBrUnconditional(loopStartLabel));
-
-                context = context.setBasicBlockLabel(loopStartLabel.name());
+                context = context.startBasicBlock(loopStartLabel.name());
 
                 context = context.enterLoop(loopStartLabel, loopEndLabel);
                 context = generateInstructions(patioLoop.body(), context);
                 context = context.exitLoop();
 
-                context = context.setBasicBlockLabel(loopEndLabel.name());
+                context = context.startBasicBlock(loopEndLabel.name());
 
                 yield context;
             }
@@ -658,7 +656,7 @@ public class LlvmPatioCodeGen {
     ) {
         var phiPredecessors = new ArrayList<PhiPredecessor>();
         for (var branch : branches) {
-            context = context.setBasicBlockLabel(branch.startLabel.name());
+            context = context.startBasicBlock(branch.startLabel.name());
 
             context = generateInstructions(branch.body.body(), context);
 
@@ -666,11 +664,7 @@ public class LlvmPatioCodeGen {
                 .anyMatch(instruction -> instruction instanceof PatioUnreachable);
 
             if (!endIsUnreachable) {
-                context = context.addTerminatorInstruction(new LlvmBrUnconditional(
-                    branch.endLabel
-                ));
-                // TODO: only allow basic block label to be set with terminator instruction?
-                context = context.setBasicBlockLabel(branch.endLabel.name());
+                context = context.startBasicBlock(branch.endLabel.name());
                 context = context.addTerminatorInstruction(new LlvmBrUnconditional(
                     endLabel
                 ));
@@ -679,7 +673,7 @@ public class LlvmPatioCodeGen {
             }
         }
 
-        context = context.setBasicBlockLabel(endLabel.name());
+        context = context.startBasicBlock(endLabel.name());
         context = generatePhiInstructions(
             targets,
             phiPredecessors,
@@ -1128,21 +1122,28 @@ public class LlvmPatioCodeGen {
             );
         }
 
-        public FunctionContext setBasicBlockLabel(String basicBlockLabel) {
-            if (this.basicBlockLabel.isEmpty()) {
-                return new FunctionContext(
-                    this.functionNameToReturnType,
-                    this.basicBlocks,
-                    Optional.of(basicBlockLabel),
-                    this.instructions,
-                    this.patioLocalToType,
-                    this.patioLocalToOperand,
-                    this.localNameAllocator,
-                    this.loopLabels
+        /**
+         * Start a new basic block with the given label. If already in a basic
+         * block, then the existing basic block is terminated with an
+         * unconditional jump to the new basic block.
+         */
+        public FunctionContext startBasicBlock(String basicBlockLabel) {
+            var context = this.basicBlockLabel.isEmpty() && this.instructions.isEmpty()
+                ? this
+                : this.addTerminatorInstruction(
+                    new LlvmBrUnconditional(new LlvmVariableLocal(basicBlockLabel))
                 );
-            } else {
-                throw new RuntimeException("basic block label is already set");
-            }
+
+            return new FunctionContext(
+                context.functionNameToReturnType,
+                context.basicBlocks,
+                Optional.of(basicBlockLabel),
+                context.instructions,
+                context.patioLocalToType,
+                context.patioLocalToOperand,
+                context.localNameAllocator,
+                context.loopLabels
+            );
         }
 
         public LlvmType patioLocalToType(PatioLocalRef value) {