Store additional data before storing free variables

3 files changed, 22 insertions, 15 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index 6631d66c..9f5bfdaf 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -18,6 +18,7 @@ import org.shedlang.compiler.util.P;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
 import java.util.stream.Stream;
 
 import static org.shedlang.compiler.util.Iterables.sortedBy;
@@ -85,6 +86,7 @@ class PatioCodeGenClosures {
         String name,
         String functionName,
         PatioDataLayout layout,
+        BiFunction<PatioLocalRef, PatioCodeGenContext, PatioCodeGenContext> storeAdditional,
         FreeVariables freeVariables,
         Source source,
         PatioCodeGenContext context
@@ -107,6 +109,8 @@ class PatioCodeGenClosures {
             source
         ));
 
+        context = storeAdditional.apply(closurePointer, context);
+
         context = storeFreeVariables(
             freeVariables,
             closurePointer,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
index 8ce61a8c..28e1ec0f 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
@@ -427,10 +427,27 @@ public class PatioCodeGenEffects {
             return bodyContext;
         });
 
+        var effectHandlerResult = context.allocateVal(
+            PatioCodeGenLocalSpace.EFFECT_HANDLER_CURRENT,
+            EFFECT_HANDLER_TYPE
+        );
+        var effectHandler = effectHandlerResult.value();
+        context = effectHandlerResult.context();
+
+        context = getCurrentEffectHandler(effectHandler, source, context);
+
         var closurePointerResult = PatioCodeGenClosures.alloc(
             name,
             fullName,
             layout,
+            (closurePointer, allocContext) ->
+                allocContext.addInstruction(new PatioMemoryStore(
+                    closurePointer,
+                    layout,
+                    CLOSURE_EFFECT_HANDLER_OFFSET,
+                    effectHandler,
+                    source
+                )),
             freeVariables,
             source,
             context
@@ -438,22 +455,7 @@ public class PatioCodeGenEffects {
         var closurePointer = closurePointerResult.value();
         context = closurePointerResult.context();
 
-        var effectHandlerResult = context.allocateVal(
-            PatioCodeGenLocalSpace.EFFECT_HANDLER_CURRENT,
-            EFFECT_HANDLER_TYPE
-        );
-        var effectHandler = effectHandlerResult.value();
-        context = effectHandlerResult.context();
 
-        context = getCurrentEffectHandler(effectHandler, source, context);
-
-        context = context.addInstruction(new PatioMemoryStore(
-            closurePointer,
-            layout,
-            CLOSURE_EFFECT_HANDLER_OFFSET,
-            effectHandler,
-            source
-        ));
 
         context = context.addInstruction(new PatioDataPtrToInt(
             target,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
index 0e8bff4a..ffa58156 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
@@ -77,6 +77,7 @@ public class PatioCodeGenFunctions {
             name,
             fullName,
             layout,
+            (closurePointer, allocContext) -> allocContext,
             freeVariables,
             source,
             context