Always store single element for additional closure data

3 files changed, 46 insertions, 37 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index 923a2fb1..6631d66c 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.patio;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.ast.typed.FreeVariable;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.patio.ast.instructions.PatioCallPointer;
@@ -29,39 +30,40 @@ class PatioCodeGenClosures {
 
     static final List<Integer> FUNCTION_POINTER_OFFSET = List.of(0);
 
+    static List<Integer> additionalOffset(List<Integer> offset) {
+        return P.vector(1).plusAll(offset);
+    }
+
+    static final PVector<Integer> FREE_VARIABLES_OFFSET = P.vector(2);
+
     private static PatioDataLayout closureDataLayoutForCall() {
         return new PatioDataLayoutStruct(List.of(
             PatioDataLayouts.FUNCTION_POINTER
         ));
     }
 
-    record Layout(PatioDataLayout dataLayout, int freeVariablesOffset) {}
-
-    static Layout closureLayout(
-        List<PatioDataLayout> additional,
+    static PatioDataLayout closureDataLayout(
+        PatioDataLayout additional,
         FreeVariables freeVariables,
         PatioCodeGenContext context
     ) {
-        var dataLayout = PatioDataLayouts.struct(
-            P.vector(PatioDataLayouts.FUNCTION_POINTER)
-                .plusAll(additional)
-                .plus(PatioDataLayouts.struct(
-                    freeVariables.stream()
-                        .map(freeVariable -> PatioDataLayouts.fromType(
-                            context.shidLocalType(freeVariable.name())
-                        ))
-                        .toList()
-                ))
-        );
-
-        return new Layout(dataLayout, 1 + additional.size());
+        return PatioDataLayouts.struct(List.of(
+            PatioDataLayouts.FUNCTION_POINTER,
+            additional,
+            PatioDataLayouts.struct(
+                freeVariables.stream()
+                    .map(freeVariable -> PatioDataLayouts.fromType(
+                        context.shidLocalType(freeVariable.name())
+                    ))
+                    .toList()
+            )
+        ));
     }
 
     private static List<Integer> freeVariableOffset(
-        Layout layout,
         FreeVariable freeVariable
     ) {
-        return List.of(layout.freeVariablesOffset, freeVariable.index());
+        return FREE_VARIABLES_OFFSET.plus(freeVariable.index());
     }
 
     static PatioCodeGenContext defineParams(
@@ -82,7 +84,7 @@ class PatioCodeGenClosures {
     static PatioCodeGenResult<PatioLocalRef> alloc(
         String name,
         String functionName,
-        Layout layout,
+        PatioDataLayout layout,
         FreeVariables freeVariables,
         Source source,
         PatioCodeGenContext context
@@ -95,11 +97,11 @@ class PatioCodeGenClosures {
         var closurePointer = closurePointerResult.value();
         context = closurePointerResult.context();
 
-        context = PatioCodeGenAlloc.alloc(closurePointer, layout.dataLayout(), source, context);
+        context = PatioCodeGenAlloc.alloc(closurePointer, layout, source, context);
 
         context = context.addInstruction(new PatioMemoryStore(
             closurePointer,
-            layout.dataLayout(),
+            layout,
             FUNCTION_POINTER_OFFSET,
             new PatioFunctionPtr(functionName),
             source
@@ -119,7 +121,7 @@ class PatioCodeGenClosures {
     static PatioCodeGenContext storeFreeVariables(
         FreeVariables freeVariables,
         PatioLocalRef closurePointer,
-        Layout closureLayout,
+        PatioDataLayout closureLayout,
         Source source,
         PatioCodeGenContext context
     ) {
@@ -128,8 +130,8 @@ class PatioCodeGenClosures {
                 freeVariable.name(),
                 c -> c.addInstruction(new PatioMemoryStore(
                     closurePointer,
-                    closureLayout.dataLayout,
-                    freeVariableOffset(closureLayout, freeVariable),
+                    closureLayout,
+                    freeVariableOffset(freeVariable),
                     c.shidLocalNameToLocal(freeVariable.name()),
                     source
                 ))
@@ -141,7 +143,7 @@ class PatioCodeGenClosures {
     static PatioCodeGenContext loadFreeVariables(
         FreeVariables freeVariables,
         PatioLocalRef closurePointer,
-        Layout closureLayout,
+        PatioDataLayout closureLayout,
         Source source,
         PatioCodeGenContext bodyContext,
         PatioCodeGenContext bodyContextParent
@@ -157,8 +159,8 @@ class PatioCodeGenClosures {
             bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
                 freeVariableLocal,
                 closurePointer,
-                closureLayout.dataLayout,
-                freeVariableOffset(closureLayout, freeVariable),
+                closureLayout,
+                freeVariableOffset(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 c0a24507..8ce61a8c 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
@@ -34,14 +34,15 @@ public class PatioCodeGenEffects {
     private static final PatioType EFFECT_HANDLER_METHOD_POINTER_TYPE = PatioTypes.DATA_POINTER;
     private static final PatioType EFFECT_HANDLER_METHOD_TYPE = PatioCodeGenTypes.VALUE_TYPE;
 
-    private static final List<Integer> CLOSURE_EFFECT_HANDLER_OFFSET = List.of(1);
+    private static final List<Integer> CLOSURE_EFFECT_HANDLER_OFFSET
+        = PatioCodeGenClosures.additionalOffset(List.of(0));
 
-    private static PatioCodeGenClosures.Layout effectHandlerMethodDataLayout(
+    private static PatioDataLayout effectHandlerMethodDataLayout(
         FreeVariables freeVariables,
         PatioCodeGenContext context
     ) {
-        return PatioCodeGenClosures.closureLayout(
-            List.of(PatioDataLayouts.DATA_POINTER),
+        return PatioCodeGenClosures.closureDataLayout(
+            PatioDataLayouts.struct(List.of(PatioDataLayouts.DATA_POINTER)),
             freeVariables,
             context
         );
@@ -400,7 +401,7 @@ public class PatioCodeGenEffects {
             bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
                 effectHandler,
                 closurePointer,
-                layout.dataLayout(),
+                layout,
                 CLOSURE_EFFECT_HANDLER_OFFSET,
                 source
             ));
@@ -448,7 +449,7 @@ public class PatioCodeGenEffects {
 
         context = context.addInstruction(new PatioMemoryStore(
             closurePointer,
-            layout.dataLayout(),
+            layout,
             CLOSURE_EFFECT_HANDLER_OFFSET,
             effectHandler,
             source
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
index cb2b8808..0e8bff4a 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenFunctions.java
@@ -4,6 +4,8 @@ import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.patio.ast.PatioVisibility;
 import org.shedlang.compiler.patio.ast.instructions.PatioDataPtrToInt;
 import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
+import org.shedlang.compiler.patio.ast.types.PatioDataLayout;
+import org.shedlang.compiler.patio.ast.types.PatioDataLayouts;
 import org.shedlang.compiler.shid.instructions.ShidDefineFunction;
 import org.shedlang.compiler.sources.Source;
 
@@ -16,11 +18,15 @@ public class PatioCodeGenFunctions {
     private PatioCodeGenFunctions() {
     }
 
-    private static PatioCodeGenClosures.Layout functionLayout(
+    private static PatioDataLayout functionDataLayout(
         FreeVariables freeVariables,
         PatioCodeGenContext context
     ) {
-        return PatioCodeGenClosures.closureLayout(List.of(), freeVariables, context);
+        return PatioCodeGenClosures.closureDataLayout(
+            PatioDataLayouts.struct(List.of()),
+            freeVariables,
+            context
+        );
     }
 
     public static PatioCodeGenContext create(
@@ -33,7 +39,7 @@ public class PatioCodeGenFunctions {
         Source source,
         PatioCodeGenContext context
     ) {
-        var layout = functionLayout(freeVariables, context);
+        var layout = functionDataLayout(freeVariables, context);
 
         var fullNameResult = context.allocateCompilationUnitElementName(
             PatioCodeGenCompilationUnitElementSpace.FUNCTION,