Extract shared closure layout code

3 files changed, 29 insertions, 42 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index a580b36b..a21aea3e 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -34,12 +34,14 @@ class PatioCodeGenClosures {
         ));
     }
 
-    static PatioDataLayout closureDataLayout(
+    record Layout(PatioDataLayout dataLayout, int freeVariablesOffset) {}
+
+    static Layout closureLayout(
         List<PatioDataLayout> additional,
         FreeVariables freeVariables,
         PatioCodeGenContext context
     ) {
-        return PatioDataLayouts.struct(
+        var dataLayout = PatioDataLayouts.struct(
             P.vector(PatioDataLayouts.FUNCTION_POINTER)
                 .plusAll(additional)
                 .plus(PatioDataLayouts.struct(
@@ -50,13 +52,15 @@ class PatioCodeGenClosures {
                         .toList()
                 ))
         );
+
+        return new Layout(dataLayout, 1 + additional.size());
     }
 
     private static List<Integer> freeVariableOffset(
-        int freeVariablesOffset,
+        Layout layout,
         FreeVariable freeVariable
     ) {
-        return List.of(freeVariablesOffset, freeVariable.index());
+        return List.of(layout.freeVariablesOffset, freeVariable.index());
     }
 
     static PatioCodeGenContext defineParams(
@@ -77,8 +81,7 @@ class PatioCodeGenClosures {
     static PatioCodeGenContext storeFreeVariables(
         FreeVariables freeVariables,
         PatioLocalRef closurePointer,
-        int freeVariablesOffset,
-        PatioDataLayout closureLayout,
+        Layout closureLayout,
         Source source,
         PatioCodeGenContext context
     ) {
@@ -87,8 +90,8 @@ class PatioCodeGenClosures {
                 freeVariable.name(),
                 c -> c.addInstruction(new PatioMemoryStore(
                     closurePointer,
-                    closureLayout,
-                    freeVariableOffset(freeVariablesOffset, freeVariable),
+                    closureLayout.dataLayout,
+                    freeVariableOffset(closureLayout, freeVariable),
                     c.shidLocalNameToLocal(freeVariable.name()),
                     source
                 ))
@@ -100,8 +103,7 @@ class PatioCodeGenClosures {
     static PatioCodeGenContext loadFreeVariables(
         FreeVariables freeVariables,
         PatioLocalRef closurePointer,
-        int freeVariablesOffset,
-        PatioDataLayout closureLayout,
+        Layout closureLayout,
         Source source,
         PatioCodeGenContext bodyContext,
         PatioCodeGenContext bodyContextParent
@@ -117,8 +119,8 @@ class PatioCodeGenClosures {
             bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
                 freeVariableLocal,
                 closurePointer,
-                closureLayout,
-                freeVariableOffset(freeVariablesOffset, freeVariable),
+                closureLayout.dataLayout,
+                freeVariableOffset(closureLayout, freeVariable),
                 source
             ));
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
index a602c50e..acf1a770 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
@@ -37,23 +37,16 @@ public class PatioCodeGenEffects {
 
     private static final List<Integer> CLOSURE_FUNCTION_POINTER_OFFSET = List.of(0);
     private static final List<Integer> CLOSURE_EFFECT_HANDLER_OFFSET = List.of(1);
-    private static final int CLOSURE_FREE_VARIABLES_OFFSET = 2;
 
-    private static PatioDataLayout effectHandlerMethodDataLayout(
+    private static PatioCodeGenClosures.Layout effectHandlerMethodDataLayout(
         FreeVariables freeVariables,
         PatioCodeGenContext context
     ) {
-        return PatioDataLayouts.struct(List.of(
-            PatioDataLayouts.FUNCTION_POINTER,
-            PatioDataLayouts.DATA_POINTER,
-            PatioDataLayouts.struct(
-                freeVariables.stream()
-                    .map(freeVariable -> PatioDataLayouts.fromType(
-                        context.shidLocalType(freeVariable.name())
-                    ))
-                    .toList()
-            )
-        ));
+        return PatioCodeGenClosures.closureLayout(
+            List.of(PatioDataLayouts.DATA_POINTER),
+            freeVariables,
+            context
+        );
     }
 
     private static final PatioFunctionDeclaration FIND_HANDLER_METHOD_DECLARATION = new PatioFunctionDeclaration(
@@ -409,7 +402,7 @@ public class PatioCodeGenEffects {
             bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
                 effectHandler,
                 closurePointer,
-                layout,
+                layout.dataLayout(),
                 CLOSURE_EFFECT_HANDLER_OFFSET,
                 source
             ));
@@ -424,7 +417,6 @@ public class PatioCodeGenEffects {
             bodyContext = loadFreeVariables(
                 freeVariables,
                 closurePointer,
-                CLOSURE_FREE_VARIABLES_OFFSET,
                 layout,
                 source,
                 bodyContext,
@@ -444,11 +436,11 @@ public class PatioCodeGenEffects {
         var closurePointer = closurePointerResult.value();
         context = closurePointerResult.context();
 
-        context = PatioCodeGenAlloc.alloc(closurePointer, layout, source, context);
+        context = PatioCodeGenAlloc.alloc(closurePointer, layout.dataLayout(), source, context);
 
         context = context.addInstruction(new PatioMemoryStore(
             closurePointer,
-            layout,
+            layout.dataLayout(),
             CLOSURE_FUNCTION_POINTER_OFFSET,
             new PatioFunctionPtr(fullName),
             source
@@ -465,7 +457,7 @@ public class PatioCodeGenEffects {
 
         context = context.addInstruction(new PatioMemoryStore(
             closurePointer,
-            layout,
+            layout.dataLayout(),
             CLOSURE_EFFECT_HANDLER_OFFSET,
             effectHandler,
             source
@@ -474,7 +466,6 @@ public class PatioCodeGenEffects {
         context = storeFreeVariables(
             freeVariables,
             closurePointer,
-            CLOSURE_FREE_VARIABLES_OFFSET,
             layout,
             source,
             context
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
index 57cca346..c2737935 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
@@ -6,7 +6,6 @@ import org.shedlang.compiler.patio.ast.instructions.PatioDataPtrToInt;
 import org.shedlang.compiler.patio.ast.instructions.PatioMemoryStore;
 import org.shedlang.compiler.patio.ast.operands.PatioFunctionPtr;
 import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
-import org.shedlang.compiler.patio.ast.types.PatioDataLayout;
 import org.shedlang.compiler.patio.ast.types.PatioTypes;
 import org.shedlang.compiler.shid.instructions.ShidDefineFunction;
 import org.shedlang.compiler.sources.Source;
@@ -20,14 +19,11 @@ public class PatioCodeGenFunctions {
     private PatioCodeGenFunctions() {
     }
 
-    // TODO: determine from layout
-    private static final int FREE_VARIABLES_OFFSET = 1;
-
-    private static PatioDataLayout functionDataLayout(
+    private static PatioCodeGenClosures.Layout functionLayout(
         FreeVariables freeVariables,
         PatioCodeGenContext context
     ) {
-        return PatioCodeGenClosures.closureDataLayout(List.of(), freeVariables, context);
+        return PatioCodeGenClosures.closureLayout(List.of(), freeVariables, context);
     }
 
     public static PatioCodeGenContext create(
@@ -40,7 +36,7 @@ public class PatioCodeGenFunctions {
         Source source,
         PatioCodeGenContext context
     ) {
-        var layout = functionDataLayout(freeVariables, context);
+        var layout = functionLayout(freeVariables, context);
 
         var fullNameResult = context.allocateCompilationUnitElementName(
             PatioCodeGenCompilationUnitElementSpace.FUNCTION,
@@ -63,7 +59,6 @@ public class PatioCodeGenFunctions {
             bodyContext = loadFreeVariables(
                 freeVariables,
                 closurePointer,
-                FREE_VARIABLES_OFFSET,
                 layout,
                 source,
                 bodyContext,
@@ -83,11 +78,11 @@ public class PatioCodeGenFunctions {
         var closurePointer = closurePointerResult.value();
         context = closurePointerResult.context();
 
-        context = PatioCodeGenAlloc.alloc(closurePointer, layout, source, context);
+        context = PatioCodeGenAlloc.alloc(closurePointer, layout.dataLayout(), source, context);
 
         context = context.addInstruction(new PatioMemoryStore(
             closurePointer,
-            layout,
+            layout.dataLayout(),
             FUNCTION_POINTER_OFFSET,
             new PatioFunctionPtr(fullName),
             source
@@ -102,7 +97,6 @@ public class PatioCodeGenFunctions {
         context = storeFreeVariables(
             freeVariables,
             closurePointer,
-            FREE_VARIABLES_OFFSET,
             layout,
             source,
             context