Store free variables in closure

3 files changed, 57 insertions, 24 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 49351b7d..82c4250b 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -128,15 +128,13 @@ public class PatioCodeGen {
                 context.temporaryPop().context();
 
             case LocalLoadInstruction localLoad -> {
-                var localResult = context.stackIrLocalNameToLocal(localLoad.name());
-                var local = localResult.value();
-                context = localResult.context();
+                var local = context.stackIrLocalNameToLocal(localLoad.name());
 
                 yield context.temporaryPush(local);
             }
 
             case LocalStoreInstruction localStore -> {
-                var localResult = context.stackIrLocalNameToLocal(localStore.name());
+                var localResult = context.addStackIrLocal(localStore.name());
                 var local = localResult.value();
                 context = localResult.context();
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index e7bef0d0..3bc18498 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.shedlang.compiler.ast.typed.FreeVariable;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.patio.ast.instructions.*;
 import org.shedlang.compiler.patio.ast.PatioLocalRef;
@@ -38,6 +39,10 @@ class PatioCodeGenClosures {
         ));
     }
 
+    private static List<Integer> freeVariableOffset(FreeVariable freeVariable) {
+        return List.of(1, freeVariable.index());
+    }
+
     public static PatioCodeGenContext create(
         PatioLocalRef target,
         String name,
@@ -97,6 +102,16 @@ class PatioCodeGenClosures {
             source
         ));
 
+        for (var freeVariable : freeVariables) {
+            context = context.addInstruction(new PatioMemoryStore(
+                closurePointer,
+                layout,
+                freeVariableOffset(freeVariable),
+                context.stackIrLocalNameToLocal(freeVariable.name()),
+                source
+            ));
+        }
+
         return context;
     }
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 8139a022..46d8e745 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -114,29 +114,45 @@ public record PatioCodeGenContext(
             ).toResult(new PatioLocalRef(fullName));
         }
 
-        public Result<PatioLocalRef> stackIrLocalNameToLocal(String name) {
+        public Result<PatioLocalRef> addStackIrLocal(String name) {
             var context = this;
 
-            if (!this.stackIrNameToLocal.containsKey(name)) {
+            if (this.stackIrNameToLocal.containsKey(name)) {
+                // TODO: better error
+                throw new RuntimeException("Stack IR local already exists: " + name);
+            }
+
+            var fullName = NameAllocator.allocateFixedName(
+                this.localsByName::containsKey,
+                PatioCodeGenLocalSpace.LOCAL,
+                name
+            );
+
+            var localRef = new PatioLocalRef(fullName);
+
+            context = new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName.plus(
+                    fullName,
+                    new PatioLocal(fullName, PatioCodeGenTypes.VALUE_TYPE)
+                ),
                 // TODO: the mapping is redundant given we fix the mapping
-                var fullName = NameAllocator.allocateFixedName(
-                    this.localsByName::containsKey,
-                    PatioCodeGenLocalSpace.LOCAL,
-                    name
-                );
-
-                context = new FunctionContext(
-                    this.temporaryStack,
-                    this.instructions,
-                    this.localsByName.plus(
-                        fullName,
-                        new PatioLocal(fullName, PatioCodeGenTypes.VALUE_TYPE)
-                    ),
-                    this.stackIrNameToLocal.plus(name, new PatioLocalRef(fullName))
-                );
+                this.stackIrNameToLocal.plus(name, localRef)
+            );
+
+            return context.toResult(localRef);
+        }
+
+        public PatioLocalRef stackIrLocalNameToLocal(String name) {
+            var local = this.stackIrNameToLocal.get(name);
+
+            if (local == null) {
+                // TODO: better error
+                throw new RuntimeException("Unknown stack IR local: " + name);
             }
 
-            return context.toResult(context.stackIrNameToLocal.get(name));
+            return local;
         }
 
         private <T> Result<T> toResult(T value) {
@@ -205,12 +221,16 @@ public record PatioCodeGenContext(
         return withFunctionContextResult(result);
     }
 
-    public PatioCodeGenResult<PatioLocalRef> stackIrLocalNameToLocal(String name) {
-        var result = this.functionContext.stackIrLocalNameToLocal(name);
+    public PatioCodeGenResult<PatioLocalRef> addStackIrLocal(String name) {
+        var result = this.functionContext.addStackIrLocal(name);
 
         return withFunctionContextResult(result);
     }
 
+    public PatioLocalRef stackIrLocalNameToLocal(String name) {
+        return this.functionContext.stackIrLocalNameToLocal(name);
+    }
+
     public PatioCodeGenContext createFunction(
         String name,
         boolean isPrivate,