Pop exception handler in catch clause

1 files changed, 65 insertions, 51 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 9a0b2b75..7bf7839c 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -364,8 +364,8 @@ public class LlvmPatioCodeGen {
                 context = generateBranches(
                     patioIf.targets(),
                     List.of(
-                        new Branch(ifTrueStartLabel, ifTrueEndLabel, patioIf.trueBody()),
-                        new Branch(ifFalseStartLabel, ifFalseEndLabel, patioIf.falseBody())
+                        Branch.of(ifTrueStartLabel, ifTrueEndLabel, patioIf.trueBody()),
+                        Branch.of(ifFalseStartLabel, ifFalseEndLabel, patioIf.falseBody())
                     ),
                     endLabel,
                     context
@@ -577,7 +577,7 @@ public class LlvmPatioCodeGen {
                     var caseEndLabel = caseEndLabelResult.value();
                     context = caseEndLabelResult.context();
 
-                    branches.add(new Branch(
+                    branches.add(Branch.of(
                         caseStartLabel,
                         caseEndLabel,
                         patioSwitchCase.body()
@@ -589,7 +589,7 @@ public class LlvmPatioCodeGen {
                     ));
                 }
 
-                branches.add(new Branch(
+                branches.add(Branch.of(
                     switchDefaultStartLabel,
                     switchDefaultEndLabel,
                     patioSwitch.defaultBody()
@@ -625,6 +625,46 @@ public class LlvmPatioCodeGen {
             }
 
             case PatioTry try_ -> {
+                var jmpBufResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.JMP_BUF);
+                var jmpBuf = jmpBufResult.value;
+                context = jmpBufResult.context;
+
+                context = context.addInstruction(new LlvmAlloca(
+                    jmpBuf,
+                    JMP_BUF_DATA
+                ));
+
+                var previousExceptionHandlerResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TRY_PREVIOUS_EXCEPTION_HANDLER);
+                var previousExceptionHandler = previousExceptionHandlerResult.value;
+                context = previousExceptionHandlerResult.context;
+
+                context = context.addInstruction(new LlvmLoad(
+                    previousExceptionHandler,
+                    LlvmTypes.pointer(LlvmTypes.VOID),
+                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
+                ));
+
+                var popExceptionHandlerInstruction = new LlvmStore(
+                    LlvmTypes.pointer(LlvmTypes.VOID),
+                    previousExceptionHandler,
+                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
+                );
+
+                context = context.addInstruction(new LlvmStore(
+                    LlvmTypes.pointer(LlvmTypes.VOID),
+                    jmpBuf,
+                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
+                ));
+
+                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)
+                ));
+
                 var tryBodyStartLabelResult = context.allocateLocalName(
                     LlvmCodeGenLocalVariableSpace.TRY_BODY_START
                 );
@@ -644,9 +684,9 @@ public class LlvmPatioCodeGen {
                 context = tryEndLabelResult.context();
 
                 var branches = new ArrayList<Branch>();
-
                 var llvmSwitchCases = new ArrayList<LlvmSwitch.Case>();
-                branches.add(new Branch(
+
+                branches.add(Branch.of(
                     tryBodyStartLabel,
                     tryBodyEndLabel,
                     try_.body()
@@ -669,6 +709,9 @@ public class LlvmPatioCodeGen {
                     branches.add(new Branch(
                         catchBodyStartLabel,
                         catchBodyEndLabel,
+                        List.of(
+                            popExceptionHandlerInstruction
+                        ),
                         catchClause.body()
                     ));
                     var tag = context.exceptionTag(catchClause.exceptionTypeName());
@@ -679,40 +722,6 @@ public class LlvmPatioCodeGen {
                 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 previousExceptionHandlerResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TRY_PREVIOUS_EXCEPTION_HANDLER);
-                var previousExceptionHandler = previousExceptionHandlerResult.value;
-                context = previousExceptionHandlerResult.context;
-
-                context = context.addInstruction(new LlvmLoad(
-                    previousExceptionHandler,
-                    LlvmTypes.pointer(LlvmTypes.VOID),
-                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
-                ));
-
-                context = context.addInstruction(new LlvmStore(
-                    LlvmTypes.pointer(LlvmTypes.VOID),
-                    jmpBuf,
-                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
-                ));
-
-                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)
-                ));
-
                 context = context.addTerminatorInstruction(new LlvmSwitch(
                     LlvmTypes.INT_32,
                     setjmpValue,
@@ -721,11 +730,7 @@ public class LlvmPatioCodeGen {
                 ));
 
                 context = context.startBasicBlock(rethrowStartLabel);
-                context = context.addInstruction(new LlvmStore(
-                    LlvmTypes.pointer(LlvmTypes.VOID),
-                    previousExceptionHandler,
-                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
-                ));
+                context = context.addInstruction(popExceptionHandlerInstruction);
                 context = context.addInstruction(RETHROW_DECLARATION.call(
                     Optional.empty(),
                     List.of(setjmpValue)
@@ -738,11 +743,7 @@ public class LlvmPatioCodeGen {
                     context
                 );
 
-                context = context.addInstruction(new LlvmStore(
-                    LlvmTypes.pointer(LlvmTypes.VOID),
-                    previousExceptionHandler,
-                    new LlvmVariableGlobal(EXCEPTION_HANDLER_DECLARATION.name())
-                ));
+                context = context.addInstruction(popExceptionHandlerInstruction);
 
                 yield context;
             }
@@ -790,8 +791,17 @@ public class LlvmPatioCodeGen {
     record Branch(
         LlvmVariableLocal startLabel,
         LlvmVariableLocal endLabel,
+        List<LlvmInstruction> prefixInstructions,
         PatioBlock body
-    ) {}
+    ) {
+        public static Branch of(
+            LlvmVariableLocal startLabel,
+            LlvmVariableLocal endLabel,
+            PatioBlock body
+        ) {
+            return new Branch(startLabel, endLabel, List.of(), body);
+        }
+    }
 
     private static FunctionContext generateBranches(
         List<PatioLocalRef> targets,
@@ -803,6 +813,10 @@ public class LlvmPatioCodeGen {
         for (var branch : branches) {
             context = context.startBasicBlock(branch.startLabel);
 
+            for (var prefixInstruction : branch.prefixInstructions) {
+                context = context.addInstruction(prefixInstruction);
+            }
+
             context = generateInstructions(branch.body.body(), context);
 
             var endIsUnreachable = branch.body.body().stream()