Store free variables into environment

3 files changed, 52 insertions, 22 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
index 2ae8b649..f194d56f 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -96,7 +96,7 @@ public class LlvmCompiler {
                     LlvmCompilerTypes.VALUE_TYPE,
                     List.of(
                         new LlvmParam(
-                            LlvmCompilerClosures.typeClosureEnvironmentPointer(),
+                            LlvmCompilerClosures.typeClosureEnvironmentPointer(Optional.of(defineFunction.freeVariables())),
                             // TODO: allocate name properly
                             "shed__env__0"
                         )
@@ -119,7 +119,7 @@ public class LlvmCompiler {
 
                 context = context.addInstruction(new LlvmCall(
                     closurePointer,
-                    LlvmCompilerClosures.typeClosurePointer(),
+                    LlvmCompilerClosures.typeClosurePointer(Optional.of(defineFunction.freeVariables())),
                     new LlvmVariableGlobal("shed__runtime__alloc"),
                     List.of(
                         new LlvmTypedOperand(LlvmTypes.INT_64, closureSize)
@@ -134,17 +134,45 @@ public class LlvmCompiler {
 
                 context = context.addInstruction(new LlvmGetElementPtr(
                     closureFunctionPointerPointer,
-                    LlvmCompilerClosures.typeClosurePointer(),
+                    LlvmCompilerClosures.typeClosurePointer(Optional.of(defineFunction.freeVariables())),
                     closurePointer,
                     List.of(LlvmTypedOperand.arrayIndex(0), LlvmTypedOperand.structureIndex(0))
                 ));
 
                 context = context.addInstruction(new LlvmStore(
-                    LlvmCompilerClosures.typeClosureFunctionPointer(),
+                    LlvmCompilerClosures.typeClosureFunctionPointer(Optional.of(defineFunction.freeVariables())),
                     new LlvmVariableGlobal(topLevelEntityName),
                     closureFunctionPointerPointer
                 ));
 
+                var freeVariableIndex = 0;
+                // TODO: change type of free variables to be list?
+                for (var freeVariable : defineFunction.freeVariables()) {
+                    var closureEnvironmentPointerResult = context.generateVariableLocal(
+                        LlvmCompilerLocalVariableSpace.CLOSURE_ENVIRONMENT_POINTER
+                    );
+                    var closureEnvironmentPointer = closureEnvironmentPointerResult.value();
+                    context = closureEnvironmentPointerResult.context();
+
+                    context = context.addInstruction(new LlvmGetElementPtr(
+                        closureEnvironmentPointer,
+                        LlvmCompilerClosures.typeClosurePointer(Optional.of(defineFunction.freeVariables())),
+                        closurePointer,
+                        List.of(
+                            LlvmTypedOperand.arrayIndex(0),
+                            LlvmTypedOperand.structureIndex(1),
+                            LlvmTypedOperand.arrayIndex(freeVariableIndex)
+                        )
+                    ));
+
+                    context = context.addInstruction(new LlvmStore(
+                        LlvmCompilerTypes.VALUE_TYPE,
+                        context.local(freeVariable),
+                        closureEnvironmentPointer
+                    ));
+                    freeVariableIndex++;
+                }
+
                 // TODO: ptrtoint
                 context = context.push(closurePointer);
 
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
index 9ea57ec9..daaa1ae1 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
@@ -18,30 +18,30 @@ public class LlvmCompilerClosures {
     private LlvmCompilerClosures() {
     }
 
-    public static LlvmTypePointer typeClosurePointer() {
-        return LlvmTypes.pointer(typeClosureStruct(Optional.empty()));
+    public static LlvmTypePointer typeClosurePointer(Optional<Set<String>> freeVariables) {
+        return LlvmTypes.pointer(typeClosureStruct(freeVariables));
     }
 
     public static LlvmTypeStructure typeClosureStruct(Optional<Set<String>> freeVariables) {
         return LlvmTypes.structure(List.of(
-            typeClosureFunctionPointer(),
+            typeClosureFunctionPointer(freeVariables),
             typeClosureEnvironment(freeVariables)
         ));
     }
 
-    public static LlvmTypePointer typeClosureFunctionPointer() {
-        return LlvmTypes.pointer(typeClosureFunction());
+    public static LlvmTypePointer typeClosureFunctionPointer(Optional<Set<String>> freeVariables) {
+        return LlvmTypes.pointer(typeClosureFunction(freeVariables));
     }
 
-    public static LlvmTypeFunction typeClosureFunction() {
+    public static LlvmTypeFunction typeClosureFunction(Optional<Set<String>> freeVariables) {
         return LlvmTypes.function(
             LlvmCompilerTypes.VALUE_TYPE,
-            List.of(typeClosureEnvironmentPointer())
+            List.of(typeClosureEnvironmentPointer(freeVariables))
         );
     }
 
-    public static LlvmTypePointer typeClosureEnvironmentPointer() {
-        return LlvmTypes.pointer(typeClosureEnvironment(Optional.empty()));
+    public static LlvmTypePointer typeClosureEnvironmentPointer(Optional<Set<String>> freeVariables) {
+        return LlvmTypes.pointer(typeClosureEnvironment(freeVariables));
     }
 
     public static LlvmTypeArray typeClosureEnvironment(Optional<Set<String>> freeVariables) {
@@ -72,7 +72,7 @@ public class LlvmCompilerClosures {
         context = callTargetResult.context();
 
         var fullArgs = new ArrayList<LlvmTypedOperand>();
-        fullArgs.add(new LlvmTypedOperand(typeClosureEnvironmentPointer(), environmentPointer));
+        fullArgs.add(new LlvmTypedOperand(typeClosureEnvironmentPointer(Optional.empty()), environmentPointer));
         for (var arg : args) {
             fullArgs.add(new LlvmTypedOperand(LlvmCompilerTypes.VALUE_TYPE, arg));
         }
@@ -99,7 +99,7 @@ public class LlvmCompilerClosures {
             closurePointer,
             LlvmCompilerTypes.VALUE_TYPE,
             value,
-            typeClosurePointer()
+            typeClosurePointer(Optional.empty())
         ));
 
         return context.toResult(closurePointer);
@@ -121,7 +121,7 @@ public class LlvmCompilerClosures {
 
         return context.addInstruction(new LlvmLoad(
             functionPointer,
-            typeClosureFunctionPointer(),
+            typeClosureFunctionPointer(Optional.empty()),
             functionPointerPointer
         )).toResult(functionPointer);
     }
@@ -138,7 +138,7 @@ public class LlvmCompilerClosures {
 
         context = context.addInstruction(new LlvmGetElementPtr(
             functionPointerPointer,
-            typeClosurePointer(),
+            typeClosurePointer(Optional.empty()),
             closurePointer,
             List.of(LlvmTypedOperand.arrayIndex(0), LlvmTypedOperand.structureIndex(0))
         ));
@@ -158,7 +158,7 @@ public class LlvmCompilerClosures {
 
         context = context.addInstruction(new LlvmGetElementPtr(
             environmentPointer,
-            typeClosurePointer(),
+            typeClosurePointer(Optional.empty()),
             closurePointer,
             List.of(LlvmTypedOperand.arrayIndex(0), LlvmTypedOperand.structureIndex(1))
         ));
diff --git a/src/main/java/org/shedlang/compiler/util/P.java b/src/main/java/org/shedlang/compiler/util/P.java
index fba3c512..5985a767 100644
--- a/src/main/java/org/shedlang/compiler/util/P.java
+++ b/src/main/java/org/shedlang/compiler/util/P.java
@@ -1,9 +1,6 @@
 package org.shedlang.compiler.util;
 
-import org.pcollections.HashTreePMap;
-import org.pcollections.PMap;
-import org.pcollections.PVector;
-import org.pcollections.TreePVector;
+import org.pcollections.*;
 
 import java.util.Arrays;
 import java.util.Map;
@@ -31,6 +28,11 @@ public class P {
     }
 
     @SafeVarargs
+    public static <T> OrderedPSet<T> orderedSet(T... values) {
+        return OrderedPSet.from(Arrays.asList(values));
+    }
+
+    @SafeVarargs
     public static <T> PVector<T> vector(T... values) {
         return TreePVector.from(Arrays.asList(values));
     }