Load free variables in closure

5 files changed, 155 insertions, 15 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 82c4250b..ffa8a74e 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -72,7 +72,7 @@ public class PatioCodeGen {
 
                 var receiverResult = context.allocateLocal(
                     PatioCodeGenLocalSpace.RECEIVER,
-                    PatioCodeGenClosures.CLOSURE_TYPE
+                    PatioCodeGenClosures.CLOSURE_POINTER_TYPE
                 );
                 var receiver = receiverResult.value();
                 context = receiverResult.context();
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index 3bc18498..97dcdd64 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -20,7 +20,7 @@ import java.util.function.Function;
 import static org.shedlang.compiler.util.Iterables.sortedBy;
 
 class PatioCodeGenClosures {
-    public static final PatioType CLOSURE_TYPE = PatioTypes.DATA_POINTER;
+    public static final PatioType CLOSURE_POINTER_TYPE = PatioTypes.DATA_POINTER;
 
     private PatioCodeGenClosures() {}
 
@@ -63,6 +63,29 @@ class PatioCodeGenClosures {
         context = fullNameResult.context();
 
         context = context.createFunction(fullName, true, source, bodyContext -> {
+            var closurePointerResult = bodyContext.allocateLocal(
+                PatioCodeGenLocalSpace.CLOSURE_POINTER,
+                CLOSURE_POINTER_TYPE
+            );
+            var closurePointer = closurePointerResult.value();
+            bodyContext = closurePointerResult.context();
+
+            bodyContext = bodyContext.addArg(closurePointer);
+
+            for (var freeVariable : freeVariables) {
+                var freeVariableLocalResult = bodyContext.addStackIrLocal(freeVariable.name());
+                var freeVariableLocal = freeVariableLocalResult.value();
+                bodyContext = freeVariableLocalResult.context();
+
+                bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
+                    freeVariableLocal,
+                    closurePointer,
+                    layout,
+                    freeVariableOffset(freeVariable),
+                    source
+                ));
+            }
+
             return bodyContext;
         });
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 46d8e745..c8879c45 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -15,6 +15,7 @@ import org.shedlang.compiler.patio.ast.types.PatioType;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.util.P;
 
+import java.util.ArrayList;
 import java.util.Optional;
 import java.util.function.Function;
 
@@ -59,13 +60,15 @@ public record PatioCodeGenContext(
         PStack<PatioLocalRef> temporaryStack,
         PVector<PatioInstruction> instructions,
         OrderedPMap<String, PatioLocal> localsByName,
-        PMap<String, PatioLocalRef> stackIrNameToLocal
+        PMap<String, PatioLocalRef> stackIrNameToLocal,
+        PVector<PatioLocalRef> args
     ) {
         public static final FunctionContext EMPTY = new FunctionContext(
             P.stack(),
             P.vector(),
             OrderedPMap.empty(),
-            P.map()
+            P.map(),
+            P.vector()
         );
 
         public Result<PatioLocalRef> temporaryPop() {
@@ -73,7 +76,8 @@ public record PatioCodeGenContext(
                 this.temporaryStack.minus(0),
                 this.instructions,
                 this.localsByName,
-                this.stackIrNameToLocal
+                this.stackIrNameToLocal,
+                this.args
             ).toResult(this.temporaryStack.get(0));
         }
 
@@ -82,7 +86,8 @@ public record PatioCodeGenContext(
                 this.temporaryStack.plus(value),
                 this.instructions,
                 this.localsByName,
-                this.stackIrNameToLocal
+                this.stackIrNameToLocal,
+                this.args
             );
         }
 
@@ -91,7 +96,8 @@ public record PatioCodeGenContext(
                 this.temporaryStack,
                 this.instructions.plus(instruction),
                 this.localsByName,
-                this.stackIrNameToLocal
+                this.stackIrNameToLocal,
+                this.args
             );
         }
 
@@ -110,7 +116,8 @@ public record PatioCodeGenContext(
                 this.temporaryStack,
                 this.instructions,
                 this.localsByName.plus(fullName, new PatioLocal(fullName, type)),
-                this.stackIrNameToLocal
+                this.stackIrNameToLocal,
+                this.args
             ).toResult(new PatioLocalRef(fullName));
         }
 
@@ -138,7 +145,8 @@ public record PatioCodeGenContext(
                     new PatioLocal(fullName, PatioCodeGenTypes.VALUE_TYPE)
                 ),
                 // TODO: the mapping is redundant given we fix the mapping
-                this.stackIrNameToLocal.plus(name, localRef)
+                this.stackIrNameToLocal.plus(name, localRef),
+                this.args
             );
 
             return context.toResult(localRef);
@@ -155,6 +163,16 @@ public record PatioCodeGenContext(
             return local;
         }
 
+        public FunctionContext addArg(PatioLocalRef local) {
+            return new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName,
+                this.stackIrNameToLocal,
+                this.args.plus(local)
+            );
+        }
+
         private <T> Result<T> toResult(T value) {
             return new Result<>(value, this);
         }
@@ -231,14 +249,43 @@ public record PatioCodeGenContext(
         return this.functionContext.stackIrLocalNameToLocal(name);
     }
 
+    public PatioCodeGenContext addArg(PatioLocalRef local) {
+        return withFunctionContext(this.functionContext.addArg(local));
+    }
+
     public PatioCodeGenContext createFunction(
         String name,
         boolean isPrivate,
         Source source,
         Function<PatioCodeGenContext, PatioCodeGenContext> createBody
     ) {
+        var bodyContext = new PatioCodeGenContext(
+            this.compilationUnitContext,
+            FunctionContext.EMPTY
+        );
+
+        bodyContext = createBody.apply(bodyContext);
+
+        var params = new ArrayList<PatioLocal>();
+        var locals = new ArrayList<PatioLocal>();
+
+        for (var local : bodyContext.functionContext.localsByName.values()) {
+            if (bodyContext.functionContext.args.contains(new PatioLocalRef(local.name()))) {
+                params.add(local);
+            } else {
+                locals.add(local);
+            }
+        }
+
         return new PatioCodeGenContext(
-            this.compilationUnitContext.addElement(new PatioFunctionDefinition(name, isPrivate, source)),
+            bodyContext.compilationUnitContext.addElement(new PatioFunctionDefinition(
+                name,
+                isPrivate,
+                params,
+                locals,
+                bodyContext.functionContext.instructions,
+                source
+            )),
             this.functionContext
         );
     }
diff --git a/src/main/java/org/shedlang/compiler/patio/ast/PatioFunctionDefinition.java b/src/main/java/org/shedlang/compiler/patio/ast/PatioFunctionDefinition.java
index d421c290..1057b2fd 100644
--- a/src/main/java/org/shedlang/compiler/patio/ast/PatioFunctionDefinition.java
+++ b/src/main/java/org/shedlang/compiler/patio/ast/PatioFunctionDefinition.java
@@ -1,29 +1,85 @@
 package org.shedlang.compiler.patio.ast;
 
+import org.pcollections.PVector;
+import org.shedlang.compiler.patio.ast.instructions.PatioInstruction;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
+import org.shedlang.compiler.util.P;
+
+import java.util.List;
 
 public record PatioFunctionDefinition(
     String name,
     boolean isPrivate,
+    List<PatioLocal> params,
+    List<PatioLocal> locals,
+    List<PatioInstruction> body,
     Source source
 ) implements PatioCompilationUnitElement {
     public static Builder builder() {
-        return new Builder("f", true);
+        return new Builder(
+            "f",
+            true,
+            P.vector(),
+            P.vector(),
+            P.vector()
+        );
     }
 
-    public record Builder(String name, boolean isPrivate) {
+    public record Builder(
+        String name,
+        boolean isPrivate,
+        PVector<PatioLocal> params,
+        PVector<PatioLocal> locals,
+        PVector<PatioInstruction> body
+    ) {
         public Builder withName(String name) {
             return new Builder(
                 name,
-                this.isPrivate
+                this.isPrivate,
+                this.params,
+                this.locals,
+                this.body
             );
         }
 
         public Builder withIsPrivate(boolean isPrivate) {
             return new Builder(
                 this.name,
-                isPrivate
+                isPrivate,
+                this.params,
+                this.locals,
+                this.body
+            );
+        }
+
+        public Builder addParam(PatioLocal param) {
+            return new Builder(
+                this.name,
+                this.isPrivate,
+                this.params.plus(param),
+                this.locals,
+                this.body
+            );
+        }
+
+        public Builder addLocal(PatioLocal local) {
+            return new Builder(
+                this.name,
+                this.isPrivate,
+                this.params,
+                this.locals.plus(local),
+                this.body
+            );
+        }
+
+        public Builder addBodyInstruction(PatioInstruction instruction) {
+            return new Builder(
+                this.name,
+                this.isPrivate,
+                this.params,
+                this.locals,
+                this.body.plus(instruction)
             );
         }
 
@@ -31,6 +87,9 @@ public record PatioFunctionDefinition(
             return new PatioFunctionDefinition(
                 this.name,
                 this.isPrivate,
+                this.params,
+                this.locals,
+                this.body,
                 Sources.NULL
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java b/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java
index 946c6eb6..cf52f9f4 100644
--- a/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java
+++ b/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java
@@ -54,7 +54,18 @@ public class PatioSerializer {
                 treeListIndented(
                     treeIdentifier("function"),
                     treeString(element.name()),
-                    treeNamed("is-private", treeBool(functionDefinition.isPrivate()))
+                    treeNamed("is-private", treeBool(functionDefinition.isPrivate())),
+                    treeNamed("params", treeList(
+                        functionDefinition.params().stream()
+                            .map(param -> localToTree(param))
+                            .toList()
+                    )),
+                    treeNamed("locals", treeList(
+                        functionDefinition.locals().stream()
+                            .map(local -> localToTree(local))
+                            .toList()
+                    )),
+                    treeNamed("body", instructionsToTree(functionDefinition.body()))
                 );
         };
     }