Support simple try/throw in LLVM backend

3 files changed, 211 insertions, 17 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 a6bd4ba4..2bd93582 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
@@ -16,6 +16,7 @@ public enum LlvmCodeGenLocalVariableSpace {
     IF_FALSE_END("ifFalseEnd"),
     IF_TRUE_START("ifTrueStart"),
     IF_TRUE_END("ifTrueEnd"),
+    JMP_BUF("jmp_buf"),
     LOCAL("local"),
     LOOP_START("loopStart"),
     LOOP_END("loopEnd"),
@@ -23,6 +24,7 @@ public enum LlvmCodeGenLocalVariableSpace {
     NAMESPACE_STRUCT_POINTER("namespaceStructPointer"),
     PARAM("param"),
     PHI("phi"),
+    SETJMP_VALUE("setjmpValue"),
     SIZE_PTR("sizePtr"),
     SIZE("size"),
     SWITCH_CASE_START("switchCaseStart"),
@@ -32,6 +34,9 @@ public enum LlvmCodeGenLocalVariableSpace {
     SWITCH_END("switchEnd"),
     TRY_BODY_START("tryBodyStart"),
     TRY_BODY_END("tryBodyEnd"),
+    TRY_CATCH_START("tryCatchStart"),
+    TRY_CATCH_END("tryCatchEnd"),
+    TRY_RETHROW_START("tryRethrowStart"),
     TRY_END("tryEnd"),
     TEMPORARY("tmp");
 
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 f1441637..2f1ce44b 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -23,8 +23,52 @@ import java.util.Optional;
 import java.util.stream.StreamSupport;
 
 public class LlvmPatioCodeGen {
+    private static final LlvmType EXCEPTION_TAG_TYPE = LlvmTypes.INT_32;
+    private static final LlvmType EXCEPTION_VALUE_TYPE = LlvmTypes.INT_64;
+
+    // TODO: find size of jmp_buf properly
+    private static final LlvmType JMP_BUF_DATA = LlvmTypes.array(200, LlvmTypes.INT_8);
+
+    private static final LlvmFunctionDeclaration SETJMP_DECLARATION = new LlvmFunctionDeclaration(
+        "_setjmp",
+        LlvmTypes.INT_32,
+        List.of(
+            new LlvmParam(LlvmTypes.pointer(LlvmTypes.VOID), "env")
+        )
+    );
+
+    private static final LlvmGlobalDefinition EXCEPTION_HANDLER_DECLARATION = new LlvmGlobalDefinition(
+        "patio__runtime__exceptions_exception_handler",
+        LlvmTypes.pointer(LlvmTypes.VOID),
+        Optional.empty(),
+        false,
+        false,
+        false
+    );
+
+    private static final LlvmFunctionDeclaration THROW_DECLARATION = new LlvmFunctionDeclaration(
+        "patio__runtime__exceptions_throw",
+        LlvmTypes.pointer(LlvmTypes.VOID),
+        List.of(
+            new LlvmParam(EXCEPTION_TAG_TYPE, "exception_tag"),
+            new LlvmParam(EXCEPTION_VALUE_TYPE, "exception_value")
+        )
+    );
+
+    private static final LlvmFunctionDeclaration RETHROW_DECLARATION = new LlvmFunctionDeclaration(
+        "patio__runtime__exceptions_rethrow",
+        LlvmTypes.pointer(LlvmTypes.VOID),
+        List.of(
+            new LlvmParam(EXCEPTION_TAG_TYPE, "exception_tag")
+        )
+    );
+
     public static LlvmModule generateCompilationUnit(PatioCompilationUnit patioCompilationUnit) {
-        var context = ModuleContext.EMPTY;
+        var context = ModuleContext.EMPTY
+            .addTopLevelEntity(SETJMP_DECLARATION)
+            .addTopLevelEntity(EXCEPTION_HANDLER_DECLARATION)
+            .addTopLevelEntity(THROW_DECLARATION)
+            .addTopLevelEntity(RETHROW_DECLARATION);
 
         for (var patioCompilationUnitElement : patioCompilationUnit.elements()) {
             switch (patioCompilationUnitElement) {
@@ -39,6 +83,10 @@ public class LlvmPatioCodeGen {
                     );
                 }
                 case PatioExceptionTypeDefinition patioExceptionTypeDefinition -> {
+                    // TODO: this needs to be done across all compilation units to ensure
+                    // consistency in tags and therefore allow cross-compilation unit throws.
+
+                    context = context.allocateExceptionTag(patioExceptionTypeDefinition.name());
                 }
                 case PatioFunctionDefinition patioFunctionDefinition -> {
                     context = context.addFunctionNameToReturnType(
@@ -564,8 +612,17 @@ public class LlvmPatioCodeGen {
                 yield context;
             }
 
-            case PatioThrow throw_ ->
-                throw new UnsupportedOperationException("TODO");
+            case PatioThrow throw_ -> {
+                context = context.addInstruction(THROW_DECLARATION.call(
+                    Optional.empty(),
+                    List.of(
+                        new LlvmOperandInt(context.exceptionTag(throw_.exceptionTypeName())),
+                        patioOperandToLlvmOperand(throw_.value(), context)
+                    )
+                ));
+
+                yield context;
+            }
 
             case PatioTry try_ -> {
                 var tryBodyStartLabelResult = context.allocateLocalName(
@@ -586,15 +643,83 @@ public class LlvmPatioCodeGen {
                 var tryEndLabel = tryEndLabelResult.value();
                 context = tryEndLabelResult.context();
 
+                var branches = new ArrayList<Branch>();
+
+                var llvmSwitchCases = new ArrayList<LlvmSwitch.Case>();
+                branches.add(new Branch(
+                    tryBodyStartLabel,
+                    tryBodyEndLabel,
+                    try_.body()
+                ));
+                llvmSwitchCases.add(new LlvmSwitch.Case(0, tryBodyStartLabel));
+
+                for (var catchClause : try_.catchClauses()) {
+                    var catchBodyStartLabelResult = context.allocateLocalName(
+                        LlvmCodeGenLocalVariableSpace.TRY_CATCH_START
+                    );
+                    var catchBodyStartLabel = catchBodyStartLabelResult.value();
+                    context = catchBodyStartLabelResult.context();
+
+                    var catchBodyEndLabelResult = context.allocateLocalName(
+                        LlvmCodeGenLocalVariableSpace.TRY_CATCH_END
+                    );
+                    var catchBodyEndLabel = catchBodyEndLabelResult.value();
+                    context = catchBodyEndLabelResult.context();
+
+                    branches.add(new Branch(
+                        catchBodyStartLabel,
+                        catchBodyEndLabel,
+                        catchClause.body()
+                    ));
+                    var tag = context.exceptionTag(catchClause.exceptionTypeName());
+                    llvmSwitchCases.add(new LlvmSwitch.Case(tag, catchBodyStartLabel));
+                }
+
+                var rethrowStartLabelResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TRY_RETHROW_START);
+                var rethrowStartLabel = rethrowStartLabelResult.value;
+                context = rethrowStartLabelResult.context;
+
+                var jmpBufResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.JMP_BUF);
+                var jmpBuf = jmpBufResult.value;
+                context = jmpBufResult.context;
+
+                context = context.addInstruction(new LlvmAlloca(
+                    jmpBuf,
+                    JMP_BUF_DATA
+                ));
+
+                var setjmpValueResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.SETJMP_VALUE);
+                var setjmpValue = setjmpValueResult.value;
+                context = setjmpValueResult.context;
+
+                context = context.addInstruction(SETJMP_DECLARATION.call(
+                    Optional.of(setjmpValue),
+                    List.of(jmpBuf)
+                ));
+
+                // TODO: save and restore previous exception handler
+                context = context.addInstruction(new LlvmStore(
+                    LlvmTypes.pointer(LlvmTypes.VOID),
+                    jmpBuf,
+                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
+                ));
+
+                context = context.addTerminatorInstruction(new LlvmSwitch(
+                    LlvmTypes.INT_32,
+                    setjmpValue,
+                    rethrowStartLabel,
+                    llvmSwitchCases
+                ));
+
+                context = context.startBasicBlock(rethrowStartLabel);
+                context = context.addInstruction(RETHROW_DECLARATION.call(
+                    Optional.empty(),
+                    List.of(setjmpValue)
+                ));
+
                 context = generateBranches(
                     try_.targets(),
-                    List.of(
-                        new Branch(
-                            tryBodyStartLabel,
-                            tryBodyEndLabel,
-                            try_.body()
-                        )
-                    ),
+                    branches,
                     tryEndLabel,
                     context
                 );
@@ -1021,24 +1146,29 @@ public class LlvmPatioCodeGen {
     private static class ModuleContext {
         private static final ModuleContext EMPTY = new ModuleContext(
             P.map(),
-            P.vector()
+            P.vector(),
+            P.map()
         );
 
         private final PMap<String, LlvmType> functionNameToReturnType;
         private final PVector<LlvmTopLevelEntity> topLevelEntities;
+        private final PMap<String, Integer> exceptionTypeNameToTag;
 
         private ModuleContext(
             PMap<String, LlvmType> functionNameToReturnType,
-            PVector<LlvmTopLevelEntity> topLevelEntities
+            PVector<LlvmTopLevelEntity> topLevelEntities,
+            PMap<String, Integer> exceptionTypeNameToTag
         ) {
             this.functionNameToReturnType = functionNameToReturnType;
             this.topLevelEntities = topLevelEntities;
+            this.exceptionTypeNameToTag = exceptionTypeNameToTag;
         }
 
         private ModuleContext addFunctionNameToReturnType(String functionName, LlvmType returnType) {
             return new ModuleContext(
                 this.functionNameToReturnType.plus(functionName, returnType),
-                this.topLevelEntities
+                this.topLevelEntities,
+                this.exceptionTypeNameToTag
             );
         }
 
@@ -1049,19 +1179,41 @@ public class LlvmPatioCodeGen {
         public ModuleContext addTopLevelEntity(LlvmTopLevelEntity topLevelEntity) {
             return new ModuleContext(
                 this.functionNameToReturnType,
-                this.topLevelEntities.plus(topLevelEntity)
+                this.topLevelEntities.plus(topLevelEntity),
+                this.exceptionTypeNameToTag
+            );
+        }
+
+        public ModuleContext allocateExceptionTag(String exceptionTypeName) {
+            var tag = this.exceptionTypeNameToTag.size() + 1;
+            return new ModuleContext(
+                this.functionNameToReturnType,
+                this.topLevelEntities,
+                this.exceptionTypeNameToTag.plus(exceptionTypeName, tag)
             );
         }
 
         public FunctionContext enterFunction() {
-            return FunctionContext.of(functionNameToReturnType);
+            return FunctionContext.of(functionNameToReturnType, exceptionTypeNameToTag);
+        }
+
+        private <T> Result<T> toResult(T value) {
+            return new Result<>(value, this);
+        }
+
+        public record Result<T>(T value, ModuleContext context) {
         }
     }
 
     private static class FunctionContext {
-        private static FunctionContext of(PMap<String, LlvmType> functionNameToType) {
+        // TODO: just accept module context?
+        private static FunctionContext of(
+            PMap<String, LlvmType> functionNameToType,
+            PMap<String, Integer> exceptionTypeNameToTag
+        ) {
             return new FunctionContext(
                 functionNameToType,
+                exceptionTypeNameToTag,
                 P.vector(),
                 Optional.empty(),
                 P.vector(),
@@ -1073,6 +1225,7 @@ public class LlvmPatioCodeGen {
         }
 
         private final PMap<String, LlvmType> functionNameToReturnType;
+        private final PMap<String, Integer> exceptionTypeNameToTag;
         private final PVector<LlvmBasicBlock> basicBlocks;
         private final Optional<String> basicBlockLabel;
         private final PVector<LlvmInstruction> instructions;
@@ -1083,7 +1236,7 @@ public class LlvmPatioCodeGen {
 
         private FunctionContext(
             PMap<String, LlvmType> functionNameToReturnType,
-            PVector<LlvmBasicBlock> basicBlocks,
+            PMap<String, Integer> exceptionTypeNameToTag, PVector<LlvmBasicBlock> basicBlocks,
             Optional<String> basicBlockLabel,
             PVector<LlvmInstruction> instructions,
             PMap<PatioLocalRef, LlvmType> patioLocalToType,
@@ -1092,6 +1245,7 @@ public class LlvmPatioCodeGen {
             PStack<LoopLabels> loopLabels
         ) {
             this.functionNameToReturnType = functionNameToReturnType;
+            this.exceptionTypeNameToTag = exceptionTypeNameToTag;
             this.basicBlocks = basicBlocks;
             this.basicBlockLabel = basicBlockLabel;
             this.instructions = instructions;
@@ -1105,9 +1259,14 @@ public class LlvmPatioCodeGen {
             return this.functionNameToReturnType.get(functionName);
         }
 
+        public int exceptionTag(String exceptionTypeName) {
+            return exceptionTypeNameToTag.get(exceptionTypeName);
+        }
+
         private FunctionContext addTerminatorInstruction(LlvmTerminatorInstruction terminatorInstruction) {
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks.plus(new LlvmBasicBlock(
                     this.basicBlockLabel,
                     this.instructions,
@@ -1136,6 +1295,7 @@ public class LlvmPatioCodeGen {
 
             return new FunctionContext(
                 context.functionNameToReturnType,
+                context.exceptionTypeNameToTag,
                 context.basicBlocks,
                 Optional.of(basicBlockLabel.name()),
                 context.instructions,
@@ -1166,6 +1326,7 @@ public class LlvmPatioCodeGen {
         public FunctionContext addPatioLocalToType(PatioLocalRef patioLocal, LlvmType llvmType) {
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1179,6 +1340,7 @@ public class LlvmPatioCodeGen {
         public FunctionContext addInstruction(LlvmInstruction instruction) {
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions.plus(instruction),
@@ -1194,6 +1356,7 @@ public class LlvmPatioCodeGen {
 
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1209,6 +1372,7 @@ public class LlvmPatioCodeGen {
 
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1228,6 +1392,7 @@ public class LlvmPatioCodeGen {
         public FunctionContext enterLoop(LlvmVariableLocal startLabel, LlvmVariableLocal endLabel) {
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1241,6 +1406,7 @@ public class LlvmPatioCodeGen {
         public FunctionContext exitLoop() {
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1262,6 +1428,7 @@ public class LlvmPatioCodeGen {
         public FunctionContext addPatioLocal(PatioLocalRef patioLocal, LlvmOperand llvmOperand) {
             return new FunctionContext(
                 this.functionNameToReturnType,
+                this.exceptionTypeNameToTag,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmFunctionDeclaration.java b/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmFunctionDeclaration.java
index 4abd406c..3987eb43 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmFunctionDeclaration.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmFunctionDeclaration.java
@@ -1,12 +1,34 @@
 package org.shedlang.compiler.backends.llvm.ir.ast;
 
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmCall;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmTypedOperand;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariable;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableGlobal;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmType;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 public record LlvmFunctionDeclaration(
     String name,
     LlvmType returnType,
     List<LlvmParam> params
 ) implements LlvmTopLevelEntity {
+    public LlvmCall call(Optional<LlvmVariable> target, List<LlvmOperand> args) {
+        var typedArgs = new ArrayList<LlvmTypedOperand>();
+
+        for (var argIndex = 0; argIndex < args.size(); argIndex++) {
+            var paramType = params.get(argIndex).type();
+            typedArgs.add(new LlvmTypedOperand(paramType, args.get(argIndex)));
+        }
+
+        return new LlvmCall(
+            target,
+            returnType,
+            new LlvmVariableGlobal(name),
+            typedArgs
+        );
+    }
 }