Extract LlvmPatioCodeGen.generateBranches()

1 files changed, 76 insertions, 53 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 b8549eae..79f2e99f 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -18,7 +18,6 @@ import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
 import java.util.ArrayList;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.StreamSupport;
@@ -298,27 +297,13 @@ public class LlvmPatioCodeGen {
                     ifFalseStartLabel.name()
                 ));
 
-                context = context.setBasicBlockLabel(ifTrueStartLabel.name());
-                context = generateInstructions(patioIf.trueBody().body(), context);
-                // TODO: only allow basic block label to be set with terminator instruction?
-                context = context.addTerminatorInstruction(new LlvmBrUnconditional(ifTrueEndLabel.name()));
-                context = context.setBasicBlockLabel(ifTrueEndLabel.name());
-                context = context.addTerminatorInstruction(new LlvmBrUnconditional(endLabel.name()));
-
-                context = context.setBasicBlockLabel(ifFalseStartLabel.name());
-                context = generateInstructions(patioIf.falseBody().body(), context);
-                context = context.addTerminatorInstruction(new LlvmBrUnconditional(ifFalseEndLabel.name()));
-                context = context.setBasicBlockLabel(ifFalseEndLabel.name());
-                context = context.addTerminatorInstruction(new LlvmBrUnconditional(endLabel.name()));
-
-                context = context.setBasicBlockLabel(endLabel.name());
-
-                context = generatePhiInstructions(
+                context = generateBranches(
                     patioIf.targets(),
                     List.of(
-                        new PhiPredecessor(patioIf.trueBody(), ifTrueEndLabel),
-                        new PhiPredecessor(patioIf.falseBody(), ifFalseEndLabel)
+                        new Branch(ifTrueStartLabel, Optional.of(ifTrueEndLabel), patioIf.trueBody()),
+                        new Branch(ifFalseStartLabel, Optional.of(ifFalseEndLabel), patioIf.falseBody())
                     ),
+                    endLabel,
                     context
                 );
 
@@ -507,7 +492,8 @@ public class LlvmPatioCodeGen {
                 var switchEndLabel = switchEndLabelResult.value();
                 context = switchEndLabelResult.context();
 
-                var switchCaseStartLabels = new LinkedHashMap<Integer, String>();
+                var llvmSwitchCases = new ArrayList<LlvmSwitch.Case>();
+                var branches = new ArrayList<Branch>();
                 for (var patioSwitchCase : patioSwitch.cases()) {
                     var caseStartLabelResult = context.allocateLocalName(
                         LlvmCodeGenLocalVariableSpace.SWITCH_CASE_START,
@@ -516,22 +502,6 @@ public class LlvmPatioCodeGen {
                     var caseStartLabel = caseStartLabelResult.value();
                     context = caseStartLabelResult.context();
 
-                    switchCaseStartLabels.put(patioSwitchCase.value(), caseStartLabel.name());
-                }
-
-                context = context.addTerminatorInstruction(new LlvmSwitch(
-                    LlvmTypes.INT_64,
-                    patioOperandToLlvmOperand(patioSwitch.value(), context),
-                    switchDefaultLabel.name(),
-                    switchCaseStartLabels.entrySet().stream()
-                        .map(entry -> new LlvmSwitch.Case(entry.getKey(), entry.getValue()))
-                        .toList()
-                ));
-
-                var phiPredecessors = new ArrayList<PhiPredecessor>();
-                for (var patioSwitchCase : patioSwitch.cases()) {
-                    var caseStartLabel = switchCaseStartLabels.get(patioSwitchCase.value());
-
                     var caseEndLabelResult = context.allocateLocalName(
                         LlvmCodeGenLocalVariableSpace.SWITCH_CASE_END,
                         Integer.toString(patioSwitchCase.value())
@@ -539,29 +509,36 @@ public class LlvmPatioCodeGen {
                     var caseEndLabel = caseEndLabelResult.value();
                     context = caseEndLabelResult.context();
 
-                    context = context.setBasicBlockLabel(caseStartLabel);
-
-                    context = generateInstructions(patioSwitchCase.body().body(), context);
-
-                    context = context.addTerminatorInstruction(new LlvmBrUnconditional(
-                        caseEndLabel.name()
-                    ));
-                    context = context.setBasicBlockLabel(caseEndLabel.name());
-                    context = context.addTerminatorInstruction(new LlvmBrUnconditional(
-                        switchEndLabel.name()
+                    branches.add(new Branch(
+                        caseStartLabel,
+                        Optional.of(caseEndLabel),
+                        patioSwitchCase.body()
                     ));
 
-                    phiPredecessors.add(new PhiPredecessor(patioSwitchCase.body(), caseEndLabel));
+                    llvmSwitchCases.add(new LlvmSwitch.Case(
+                        patioSwitchCase.value(),
+                        caseStartLabel.name()
+                    ));
                 }
 
-                // TODO: add predecessor for default destination if not unreachable
-                context = context.setBasicBlockLabel(switchDefaultLabel.name());
-                context = generateInstructions(patioSwitch.defaultBody().body(), context);
+                branches.add(new Branch(
+                    switchDefaultLabel,
+                    // TODO: handle defaults that don't terminate
+                    Optional.empty(),
+                    patioSwitch.defaultBody()
+                ));
 
-                context = context.setBasicBlockLabel(switchEndLabel.name());
-                context = generatePhiInstructions(
+                context = context.addTerminatorInstruction(new LlvmSwitch(
+                    LlvmTypes.INT_64,
+                    patioOperandToLlvmOperand(patioSwitch.value(), context),
+                    switchDefaultLabel.name(),
+                    llvmSwitchCases
+                ));
+
+                context = generateBranches(
                     patioSwitch.targets(),
-                    phiPredecessors,
+                    branches,
+                    switchEndLabel,
                     context
                 );
 
@@ -608,6 +585,52 @@ public class LlvmPatioCodeGen {
         };
     }
 
+    record Branch(
+        LlvmVariableLocal startLabel,
+        Optional<LlvmVariableLocal> endLabel,
+        PatioBlock body
+    ) {}
+
+    private static FunctionContext generateBranches(
+        List<PatioLocalRef> targets,
+        List<Branch> branches,
+        LlvmVariableLocal endLabel,
+        FunctionContext context
+    ) {
+        var phiPredecessors = new ArrayList<PhiPredecessor>();
+        for (var branch : branches) {
+            context = context.setBasicBlockLabel(branch.startLabel.name());
+
+            context = generateInstructions(branch.body.body(), context);
+
+            var endIsUnreachable = branch.body.body().stream()
+                .anyMatch(instruction -> instruction instanceof PatioUnreachable);
+
+            if (!endIsUnreachable) {
+                var branchEndLabel = branch.endLabel.orElseThrow();
+                context = context.addTerminatorInstruction(new LlvmBrUnconditional(
+                    branchEndLabel.name()
+                ));
+                // TODO: only allow basic block label to be set with terminator instruction?
+                context = context.setBasicBlockLabel(branchEndLabel.name());
+                context = context.addTerminatorInstruction(new LlvmBrUnconditional(
+                    endLabel.name()
+                ));
+
+                phiPredecessors.add(new PhiPredecessor(branch.body, branchEndLabel));
+            }
+        }
+
+        context = context.setBasicBlockLabel(endLabel.name());
+        context = generatePhiInstructions(
+            targets,
+            phiPredecessors,
+            context
+        );
+
+        return context;
+    }
+
     record PhiPredecessor(PatioBlock block, LlvmVariableLocal label) {}
 
     private static FunctionContext generatePhiInstructions(