Support code generation for ifs in LLVM backend

3 files changed, 124 insertions, 16 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
index ee95777f..f937de4e 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -7,6 +7,7 @@ import org.shedlang.compiler.backends.llvm.ir.ast.instructions.*;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperandInt;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariable;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.typechecker.TypeInfo;
@@ -61,13 +62,7 @@ public class LlvmCodeGen {
 
         context = LlvmCodeGenNamespaces.define(
             namespaceName,
-            initContext -> {
-                for (var instruction : instructions) {
-                    initContext = generateInstruction(instruction, initContext);
-                }
-
-                return initContext;
-            },
+            initContext -> generateInstructions(instructions, initContext),
             context
         );
 
@@ -120,12 +115,8 @@ public class LlvmCodeGen {
                     defineFunction.positionalParams(),
                     defineFunction.namedParams(),
                     defineFunction.freeVariables(),
-                    functionContext -> {
-                        for (var bodyInstruction : defineFunction.body()) {
-                            functionContext = generateInstruction(bodyInstruction, functionContext);
-                        }
-                        return functionContext;
-                    },
+                    functionContext ->
+                        generateInstructions(defineFunction.body(), functionContext),
                     context
                 );
             }
@@ -149,8 +140,79 @@ public class LlvmCodeGen {
                 yield context.push(fieldValue);
             }
 
-            case IfInstruction if_ ->
-                throw new UnsupportedOperationException("TODO");
+            case IfInstruction if_ -> {
+                var conditionResult = context.pop();
+                var condition = conditionResult.value();
+                context = conditionResult.context();
+
+                var condition1Result = context.generateVariableLocal(
+                    LlvmCodeGenLocalVariableSpace.CONDITION
+                );
+                var condition1 = condition1Result.value();
+                context = condition1Result.context();
+
+                var ifTrueLabelResult = context.generateVariableLocal(
+                    LlvmCodeGenLocalVariableSpace.IF_TRUE
+                );
+                var ifTrueLabel = ifTrueLabelResult.value();
+                context = ifTrueLabelResult.context();
+
+                var ifFalseLabelResult = context.generateVariableLocal(
+                    LlvmCodeGenLocalVariableSpace.IF_FALSE
+                );
+                var ifFalseLabel = ifFalseLabelResult.value();
+                context = ifFalseLabelResult.context();
+
+                var endLabelResult = context.generateVariableLocal(
+                    LlvmCodeGenLocalVariableSpace.IF_END
+                );
+                var endLabel = endLabelResult.value();
+                context = endLabelResult.context();
+
+                context = context
+                    .addInstruction(new LlvmTrunc(
+                        condition1,
+                        LlvmCodeGenTypes.VALUE_TYPE,
+                        condition,
+                        LlvmTypes.INT_1
+                    ))
+                    .addTerminatorInstruction(new LlvmBr(condition1, ifTrueLabel.name(), ifFalseLabel.name()));
+
+                context = context.setBasicBlockLabel(ifTrueLabel.name());
+                context = generateInstructions(if_.trueBody(), context);
+                var trueValueResult = context.pop();
+                var trueValue = trueValueResult.value();
+                context = trueValueResult.context();
+                // TODO: remove this hack!
+                var ifTrueEndLabel = context.functionContext().basicBlockLabel().orElseThrow();
+                context = context.addTerminatorInstruction(new LlvmBrUnconditional(endLabel.name()));
+
+                context = context.setBasicBlockLabel(ifFalseLabel.name());
+                context = generateInstructions(if_.falseBody(), context);
+                var falseValueResult = context.pop();
+                var falseValue = falseValueResult.value();
+                context = falseValueResult.context();
+                // TODO: remove this hack!
+                var ifFalseEndLabel = context.functionContext().basicBlockLabel().orElseThrow();
+                context = context.addTerminatorInstruction(new LlvmBrUnconditional(endLabel.name()));
+
+                context = context.setBasicBlockLabel(endLabel.name());
+
+                var phiResult = context.generateVariableLocal(LlvmCodeGenLocalVariableSpace.PHI);
+                var phi = phiResult.value();
+                context = phiResult.context();
+
+                context = context.addInstruction(new LlvmPhi(
+                    phi,
+                    LlvmCodeGenTypes.VALUE_TYPE,
+                    List.of(
+                        new LlvmPhiPair(trueValue, new LlvmVariableLocal(ifTrueEndLabel)),
+                        new LlvmPhiPair(falseValue, new LlvmVariableLocal(ifFalseEndLabel))
+                    )
+                )).push(phi);
+
+                yield context;
+            }
 
             case IntAddInstruction intAdd ->
                 generateBinaryOperation(
@@ -247,6 +309,16 @@ public class LlvmCodeGen {
         };
     }
 
+    private static LlvmCodeGenContext generateInstructions(
+        List<Instruction> instructions,
+        LlvmCodeGenContext functionContext
+    ) {
+        for (var bodyInstruction : instructions) {
+            functionContext = generateInstruction(bodyInstruction, functionContext);
+        }
+        return functionContext;
+    }
+
     interface BinaryOperation {
         LlvmCodeGenContext generate(
             LlvmVariable target,
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
index 90541aca..eef011f4 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
@@ -67,6 +67,7 @@ public record LlvmCodeGenContext(
         PStack<LlvmOperand> stack,
         PMap<String, LlvmOperand> shedLocalNameToOperand,
         PVector<LlvmBasicBlock> basicBlocks,
+        Optional<String> basicBlockLabel,
         PVector<LlvmInstruction> instructions,
         NameAllocator<LlvmCodeGenLocalVariableSpace> localVariableNames
     ) {
@@ -74,6 +75,7 @@ public record LlvmCodeGenContext(
             Empty.stack(),
             Empty.map(),
             Empty.vector(),
+            Optional.empty(),
             Empty.vector(),
             NameAllocator.empty()
         );
@@ -83,6 +85,7 @@ public record LlvmCodeGenContext(
                 stack.plus(operand),
                 shedLocalNameToOperand,
                 basicBlocks,
+                basicBlockLabel,
                 instructions,
                 localVariableNames
             );
@@ -93,6 +96,7 @@ public record LlvmCodeGenContext(
                 stack.minus(0),
                 shedLocalNameToOperand,
                 basicBlocks,
+                basicBlockLabel,
                 instructions,
                 localVariableNames
             ).toResult(stack.get(0));
@@ -103,6 +107,7 @@ public record LlvmCodeGenContext(
                 stack,
                 shedLocalNameToOperand.plus(name, value),
                 basicBlocks,
+                basicBlockLabel,
                 instructions,
                 localVariableNames
             );
@@ -118,11 +123,28 @@ public record LlvmCodeGenContext(
             }
         }
 
+        public FunctionContext setBasicBlockLabel(String basicBlockLabel) {
+            if (this.basicBlockLabel.isPresent()) {
+                // TODO: better error
+                throw new RuntimeException("basic block label is already set");
+            }
+
+            return new FunctionContext(
+                stack,
+                shedLocalNameToOperand,
+                basicBlocks,
+                Optional.of(basicBlockLabel),
+                instructions,
+                localVariableNames
+            );
+        }
+
         public FunctionContext addInstruction(LlvmInstruction instruction) {
             return new FunctionContext(
                 stack,
                 shedLocalNameToOperand,
                 basicBlocks,
+                basicBlockLabel,
                 instructions.plus(instruction),
                 localVariableNames
             );
@@ -135,10 +157,11 @@ public record LlvmCodeGenContext(
                 stack,
                 shedLocalNameToOperand,
                 basicBlocks.plus(new LlvmBasicBlock(
-                    Optional.empty(),
+                    basicBlockLabel,
                     instructions,
                     terminatorInstruction
                 )),
+                Optional.empty(),
                 Empty.vector(),
                 localVariableNames
             );
@@ -153,6 +176,7 @@ public record LlvmCodeGenContext(
                 stack,
                 shedLocalNameToOperand,
                 basicBlocks,
+                basicBlockLabel,
                 instructions,
                 result.allocator()
             ).toResult(new LlvmVariableLocal(result.name()));
@@ -281,6 +305,13 @@ public record LlvmCodeGenContext(
         return functionContext.basicBlocks;
     }
 
+    public LlvmCodeGenContext setBasicBlockLabel(String label) {
+        return new LlvmCodeGenContext(
+            this.moduleContext,
+            this.functionContext.setBasicBlockLabel(label)
+        );
+    }
+
     public List<LlvmInstruction> instructions() {
         return functionContext.instructions;
     }
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 3c5a3aef..1d893ff0 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
@@ -8,11 +8,16 @@ public enum LlvmCodeGenLocalVariableSpace {
     CLOSURE_ENVIRONMENT_POINTER("closureEnvironmentPointer"),
     CLOSURE_FREE_VARIABLE_POINTER("closureFreeVariablePointer"),
     CLOSURE_FREE_VARIABLE("closureFreeVariable"),
+    CONDITION("condition"),
     FIELD("field"),
     FIELD_POINTER("fieldPointer"),
+    IF_END("ifEnd"),
+    IF_FALSE("ifFalse"),
+    IF_TRUE("ifTrue"),
     NAMESPACE("namespace"),
     NAMESPACE_STRUCT_POINTER("namespaceStructPointer"),
     PARAM("param"),
+    PHI("phi"),
     SIZE_PTR("sizePtr"),
     SIZE("size"),
     TEMPORARY("tmp");