Handle returns from try bodies

3 files changed, 87 insertions, 21 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 6a18af6b..262fc2c5 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -615,6 +615,11 @@ public class PatioCodeGen {
                 var value = valueResult.value();
                 context = valueResult.context();
 
+                for (var onReturnInstruction : context.onReturn()) {
+                    context = context.addInstruction(onReturnInstruction);
+                }
+
+
                 context = PatioCodeGenClosures.beforeExit(
                     List.of(value),
                     return_.source(),
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 87d5901a..40bdf823 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -188,7 +188,8 @@ public record PatioCodeGenContext(
         PMap<String, PatioLocalRef> shidNameToLocal,
         PMap<String, PVector<Function<PatioCodeGenContext, PatioCodeGenContext>>> waitingForShidLocal,
         Optional<PatioOperand> stashedEffectHandler,
-        Optional<PatioOperand> stashedGcRootStackIndex
+        Optional<PatioOperand> stashedGcRootStackIndex,
+        PStack<PatioInstruction> onReturn
     ) {
         public static final FunctionContext EMPTY = new FunctionContext(
             P.stack(),
@@ -197,7 +198,8 @@ public record PatioCodeGenContext(
             P.map(),
             P.map(),
             Optional.empty(),
-            Optional.empty()
+            Optional.empty(),
+            P.stack()
         );
 
         @CheckReturnValue
@@ -209,7 +211,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             ).toResult(this.temporaryStack.get(0));
         }
 
@@ -222,7 +225,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             ).toResult(this.temporaryStack);
         }
 
@@ -235,7 +239,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             );
         }
 
@@ -248,7 +253,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             );
         }
 
@@ -271,7 +277,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             ).toResult(new PatioLocalRef(fullName));
         }
 
@@ -307,7 +314,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal.plus(name, localRef),
                 this.waitingForShidLocal.plus(name, P.vector()),
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             );
 
             return context.toResult(localRef);
@@ -322,7 +330,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal.minus(name),
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             );
 
             return context.toResult(this.waitingForShidLocal.getOrDefault(name, P.vector()));
@@ -373,7 +382,8 @@ public record PatioCodeGenContext(
                     this.waitingForShidLocal.getOrDefault(name, P.vector()).plus(func)
                 ),
                 this.stashedEffectHandler,
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             );
         }
 
@@ -387,7 +397,8 @@ public record PatioCodeGenContext(
                 this.waitingForShidLocal,
                 // TODO: check current value is empty
                 Optional.of(effectHandler),
-                this.stashedGcRootStackIndex
+                this.stashedGcRootStackIndex,
+                this.onReturn
             );
         }
 
@@ -408,7 +419,8 @@ public record PatioCodeGenContext(
                 this.waitingForShidLocal,
                 this.stashedEffectHandler,
                 // TODO: check current value is empty
-                Optional.of(gcRootStackIndex)
+                Optional.of(gcRootStackIndex),
+                this.onReturn
             );
         }
 
@@ -419,6 +431,32 @@ public record PatioCodeGenContext(
                 .orElseThrow(() -> new RuntimeException("No stashed GC root stack index"));
         }
 
+        public FunctionContext pushOnReturn(PatioInstruction onReturn) {
+            return new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName,
+                this.shidNameToLocal,
+                this.waitingForShidLocal,
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex,
+                this.onReturn.plus(onReturn)
+            );
+        }
+
+        public FunctionContext popOnReturn() {
+            return new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName,
+                this.shidNameToLocal,
+                this.waitingForShidLocal,
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex,
+                this.onReturn.minus(0)
+            );
+        }
+
         private <T> Result<T> toResult(T value) {
             return new Result<>(value, this);
         }
@@ -724,7 +762,8 @@ public record PatioCodeGenContext(
                 this.functionContext.shidNameToLocal,
                 this.functionContext.waitingForShidLocal,
                 this.functionContext.stashedEffectHandler,
-                this.functionContext.stashedGcRootStackIndex
+                this.functionContext.stashedGcRootStackIndex,
+                this.functionContext.onReturn
             )
         );
     }
@@ -740,13 +779,15 @@ public record PatioCodeGenContext(
             this.compilationUnitContext,
             this.blockStack.minus(0),
             new FunctionContext(
+                // TODO: some of these fields should be popped
                 this.blockStack.get(0).temporaryStack,
                 this.blockStack.get(0).instructions,
                 this.functionContext.localsByName,
                 this.functionContext.shidNameToLocal,
                 this.functionContext.waitingForShidLocal,
                 this.functionContext.stashedEffectHandler,
-                this.functionContext.stashedGcRootStackIndex
+                this.functionContext.stashedGcRootStackIndex,
+                this.blockStack.get(0).onReturn
             )
         ).toResult(this.functionContext.instructions);
     }
@@ -778,6 +819,21 @@ public record PatioCodeGenContext(
     }
 
     @CheckReturnValue
+    public PatioCodeGenContext pushOnReturn(PatioInstruction onReturn) {
+        return withFunctionContext(this.functionContext.pushOnReturn(onReturn));
+    }
+
+    @CheckReturnValue
+    public PatioCodeGenContext popOnReturn() {
+        return withFunctionContext(this.functionContext.popOnReturn());
+    }
+
+    @CheckReturnValue
+    public List<PatioInstruction> onReturn() {
+        return this.functionContext.onReturn();
+    }
+
+    @CheckReturnValue
     private PatioCodeGenContext withCompilationUnitContext(CompilationUnitContext compilationUnitContext) {
         return new PatioCodeGenContext(this.globalContext, compilationUnitContext, this.blockStack, this.functionContext);
     }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
index e72055b5..d5ef5f0e 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
@@ -326,10 +326,21 @@ public class PatioCodeGenEffects {
             ));
         }
 
+        var popHandler = new PatioCall(
+            Optional.empty(),
+            POP_HANDLER_DECLARATION.name(),
+            List.of(),
+            source
+        );
+
+        context = context.pushOnReturn(popHandler);
+
         var bodyResult = generateBody.apply(context);
         var body = bodyResult.value();
         context = bodyResult.context();
 
+        context = context.popOnReturn();
+
         var exceptionValueResult = context.allocateVal(
             PatioCodeGenLocalSpace.TEMPORARY,
             PatioCodeGenTypes.VALUE_TYPE
@@ -351,13 +362,7 @@ public class PatioCodeGenEffects {
                     source
                 )
             ),
-            // TODO: also need to pop on early exit from the body
-            body.addInstruction(new PatioCall(
-                Optional.empty(),
-                POP_HANDLER_DECLARATION.name(),
-                List.of(),
-                source
-            )),
+            body.addInstruction(popHandler),
             source
         ));