Add LLVM code gen for Patio switch

2 files changed, 116 insertions, 23 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
index 8e2967df..989c657f 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
@@ -25,6 +25,10 @@ public enum LlvmCodeGenLocalVariableSpace {
     PHI("phi"),
     SIZE_PTR("sizePtr"),
     SIZE("size"),
+    SWITCH_CASE_START("switchCaseStart"),
+    SWITCH_CASE_END("switchCaseEnd"),
+    SWITCH_DEFAULT("switchDefault"),
+    SWITCH_END("switchEnd"),
     TEMPORARY("tmp");
 
     private final String name;
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 45745796..828186e3 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -17,6 +17,8 @@ import org.shedlang.compiler.patio.ast.types.*;
 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;
@@ -311,28 +313,14 @@ public class LlvmPatioCodeGen {
 
                 context = context.setBasicBlockLabel(endLabel.name());
 
-                for (var targetIndex = 0; targetIndex < patioIf.targets().size(); targetIndex++) {
-                    var patioTarget = patioIf.targets().get(targetIndex);
-
-                    var llvmTargetResult = context.addPatioLocalToLlvmLocal(patioTarget);
-                    var llvmTarget = llvmTargetResult.value;
-                    context = llvmTargetResult.context;
-
-                    context = context.addInstruction(new LlvmPhi(
-                        llvmTarget,
-                        context.patioLocalToType(patioTarget),
-                        List.of(
-                            new LlvmPhiPair(
-                                patioOperandToLlvmOperand(patioIf.trueBody().results().get(targetIndex), context),
-                                ifTrueEndLabel
-                            ),
-                            new LlvmPhiPair(
-                                patioOperandToLlvmOperand(patioIf.falseBody().results().get(targetIndex), context),
-                                ifFalseEndLabel
-                            )
-                        )
-                    ));
-                }
+                context = generatePhiInstructions(
+                    patioIf.targets(),
+                    List.of(
+                        new PhiPredecessor(patioIf.trueBody(), ifTrueEndLabel),
+                        new PhiPredecessor(patioIf.falseBody(), ifFalseEndLabel)
+                    ),
+                    context
+                );
 
                 yield context;
             }
@@ -507,7 +495,77 @@ public class LlvmPatioCodeGen {
                 ));
 
             case PatioSwitch patioSwitch -> {
-                throw new UnsupportedOperationException("TODO");
+                var defaultLabelResult = context.allocateLocalName(
+                    LlvmCodeGenLocalVariableSpace.SWITCH_DEFAULT
+                );
+                var switchDefaultLabel = defaultLabelResult.value();
+                context = defaultLabelResult.context();
+
+                var switchEndLabelResult = context.allocateLocalName(
+                    LlvmCodeGenLocalVariableSpace.SWITCH_END
+                );
+                var switchEndLabel = switchEndLabelResult.value();
+                context = switchEndLabelResult.context();
+
+                var switchCaseStartLabels = new LinkedHashMap<Integer, String>();
+                for (var patioSwitchCase : patioSwitch.cases()) {
+                    var caseStartLabelResult = context.allocateLocalName(
+                        LlvmCodeGenLocalVariableSpace.SWITCH_CASE_START,
+                        Integer.toString(patioSwitchCase.value())
+                    );
+                    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())
+                    );
+                    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()
+                    ));
+
+                    phiPredecessors.add(new PhiPredecessor(patioSwitchCase.body(), caseEndLabel));
+                }
+
+                // TODO: add predecessor for default destination if not unreachable
+                context = context.setBasicBlockLabel(switchDefaultLabel.name());
+                context = generateInstructions(patioSwitch.defaultBody().body(), context);
+
+                context = context.setBasicBlockLabel(switchEndLabel.name());
+                context = generatePhiInstructions(
+                    patioSwitch.targets(),
+                    phiPredecessors,
+                    context
+                );
+
+                yield context;
             }
 
             case PatioUnreachable patioUnreachable ->
@@ -550,6 +608,37 @@ public class LlvmPatioCodeGen {
         };
     }
 
+    record PhiPredecessor(PatioBlock block, LlvmVariableLocal label) {}
+
+    private static FunctionContext generatePhiInstructions(
+        List<PatioLocalRef> targets,
+        List<PhiPredecessor> predecessors,
+        FunctionContext context
+    ) {
+        for (var targetIndex = 0; targetIndex < targets.size(); targetIndex++) {
+            var patioTarget = targets.get(targetIndex);
+
+            var llvmTargetResult = context.addPatioLocalToLlvmLocal(patioTarget);
+            var llvmTarget = llvmTargetResult.value;
+            context = llvmTargetResult.context;
+
+            var blockTargetIndex = targetIndex;
+            var blockContext = context;
+            context = context.addInstruction(new LlvmPhi(
+                llvmTarget,
+                context.patioLocalToType(patioTarget),
+                predecessors.stream()
+                    .map(predecessor -> new LlvmPhiPair(
+                        patioOperandToLlvmOperand(predecessor.block.results().get(blockTargetIndex), blockContext),
+                        predecessor.label
+                    ))
+                    .toList()
+            ));
+        }
+
+        return context;
+    }
+
     private static PVector<LlvmTypedOperand> patioDataOffsetsToLlvmOffsets(
         PatioDataLayout patioDataLayout,
         List<Integer> patioOffsets