Store free variables into environment

5 files changed, 118 insertions, 32 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));
     }
diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerCallTests.java b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerCallTests.java
index 428c9d5d..26804c2b 100644
--- a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerCallTests.java
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerCallTests.java
@@ -10,6 +10,7 @@ import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
 import org.shedlang.compiler.stackir.instructions.Instructions;
 
 import java.util.List;
+import java.util.Optional;
 
 import static org.shedlang.compiler.backends.llvm.ir.ast.LlvmIrMatchers.*;
 import static org.zwobble.precisely.AssertThat.assertThat;
@@ -31,12 +32,12 @@ public class LlvmCompilerCallTests {
                 has("target", x -> x.target(), isLlvmVariableLocal("closurePointer__0")),
                 has("sourceType", x -> x.sourceType(), equalTo(LlvmCompilerTypes.VALUE_TYPE)),
                 has("value", x -> x.value(), isLlvmVariableLocal("f")),
-                has("targetType", x -> x.targetType(), equalTo(LlvmCompilerClosures.typeClosurePointer()))
+                has("targetType", x -> x.targetType(), equalTo(LlvmCompilerClosures.typeClosurePointer(Optional.empty())))
             ),
             instanceOf(
                 LlvmGetElementPtr.class,
                 has("target", x -> x.target(), isLlvmVariableLocal("closureFunctionPointerPointer__0")),
-                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer())),
+                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer(Optional.empty()))),
                 has("pointer", x -> x.pointer(), isLlvmVariableLocal("closurePointer__0")),
                 has("indexes", x -> x.indexes(), isSequence(
                     isLlvmTypedOperand(LlvmTypes.INT_64, isLlvmOperandInt(0)),
@@ -46,13 +47,13 @@ public class LlvmCompilerCallTests {
             instanceOf(
                 LlvmLoad.class,
                 has("target", x -> x.target(), isLlvmVariableLocal("closureFunctionPointer__0")),
-                has("type", x -> x.type(), equalTo(LlvmCompilerClosures.typeClosureFunctionPointer())),
+                has("type", x -> x.type(), equalTo(LlvmCompilerClosures.typeClosureFunctionPointer(Optional.empty()))),
                 has("pointer", x -> x.pointer(), isLlvmVariableLocal("closureFunctionPointerPointer__0"))
             ),
             instanceOf(
                 LlvmGetElementPtr.class,
                 has("target", x -> x.target(), isLlvmVariableLocal("closureEnvironmentPointer__0")),
-                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer())),
+                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer(Optional.empty()))),
                 has("pointer", x -> x.pointer(), isLlvmVariableLocal("closurePointer__0")),
                 has("indexes", x -> x.indexes(), isSequence(
                     isLlvmTypedOperand(LlvmTypes.INT_64, isLlvmOperandInt(0)),
@@ -66,7 +67,7 @@ public class LlvmCompilerCallTests {
                 has("functionPointer", x -> x.functionPointer(), isLlvmVariableLocal("closureFunctionPointer__0")),
                 has("args", x -> x.args(), isSequence(
                     isLlvmTypedOperand(
-                        LlvmCompilerClosures.typeClosureEnvironmentPointer(),
+                        LlvmCompilerClosures.typeClosureEnvironmentPointer(Optional.empty()),
                         isLlvmVariableLocal("closureEnvironmentPointer__0")
                     )
                 ))
@@ -97,12 +98,12 @@ public class LlvmCompilerCallTests {
                 has("target", x -> x.target(), isLlvmVariableLocal("closurePointer__0")),
                 has("sourceType", x -> x.sourceType(), equalTo(LlvmCompilerTypes.VALUE_TYPE)),
                 has("value", x -> x.value(), isLlvmVariableLocal("f")),
-                has("targetType", x -> x.targetType(), equalTo(LlvmCompilerClosures.typeClosurePointer()))
+                has("targetType", x -> x.targetType(), equalTo(LlvmCompilerClosures.typeClosurePointer(Optional.empty())))
             ),
             instanceOf(
                 LlvmGetElementPtr.class,
                 has("target", x -> x.target(), isLlvmVariableLocal("closureFunctionPointerPointer__0")),
-                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer())),
+                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer(Optional.empty()))),
                 has("pointer", x -> x.pointer(), isLlvmVariableLocal("closurePointer__0")),
                 has("indexes", x -> x.indexes(), isSequence(
                     isLlvmTypedOperand(LlvmTypes.INT_64, isLlvmOperandInt(0)),
@@ -112,13 +113,13 @@ public class LlvmCompilerCallTests {
             instanceOf(
                 LlvmLoad.class,
                 has("target", x -> x.target(), isLlvmVariableLocal("closureFunctionPointer__0")),
-                has("type", x -> x.type(), equalTo(LlvmCompilerClosures.typeClosureFunctionPointer())),
+                has("type", x -> x.type(), equalTo(LlvmCompilerClosures.typeClosureFunctionPointer(Optional.empty()))),
                 has("pointer", x -> x.pointer(), isLlvmVariableLocal("closureFunctionPointerPointer__0"))
             ),
             instanceOf(
                 LlvmGetElementPtr.class,
                 has("target", x -> x.target(), isLlvmVariableLocal("closureEnvironmentPointer__0")),
-                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer())),
+                has("pointerType", x -> x.pointerType(), equalTo(LlvmCompilerClosures.typeClosurePointer(Optional.empty()))),
                 has("pointer", x -> x.pointer(), isLlvmVariableLocal("closurePointer__0")),
                 has("indexes", x -> x.indexes(), isSequence(
                     isLlvmTypedOperand(LlvmTypes.INT_64, isLlvmOperandInt(0)),
@@ -132,7 +133,7 @@ public class LlvmCompilerCallTests {
                 has("functionPointer", x -> x.functionPointer(), isLlvmVariableLocal("closureFunctionPointer__0")),
                 has("args", x -> x.args(), isSequence(
                     isLlvmTypedOperand(
-                        LlvmCompilerClosures.typeClosureEnvironmentPointer(),
+                        LlvmCompilerClosures.typeClosureEnvironmentPointer(Optional.empty()),
                         isLlvmVariableLocal("closureEnvironmentPointer__0")
                     ),
                     isLlvmTypedOperand(LlvmCompilerTypes.VALUE_TYPE, isLlvmVariableLocal("positional1")),
diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerDefineFunctionTests.java b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerDefineFunctionTests.java
index 5074cee8..40143dce 100644
--- a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerDefineFunctionTests.java
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerDefineFunctionTests.java
@@ -2,10 +2,12 @@ package org.shedlang.compiler.backends.llvm;
 
 import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.backends.llvm.ir.LlvmIrSerializer;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmIr;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.stackir.instructions.DefineFunctionInstruction;
 import org.shedlang.compiler.stackir.instructions.Instructions;
+import org.shedlang.compiler.util.P;
 import org.zwobble.precisely.Matcher;
 
 import java.util.List;
@@ -66,6 +68,59 @@ public class LlvmCompilerDefineFunctionTests {
         assertThat(result.locals().entrySet(), isSequence());
     }
 
+    @Test
+    public void freeVariablesAreStoredInAndLoadedFromEnvironment() {
+        var instruction = new DefineFunctionInstruction(
+            "f",
+            List.of(),
+            List.of(),
+            P.orderedSet("a", "b"),
+            List.of(
+                Instructions.pushUnit(),
+                Instructions.return_()
+            ),
+            Sources.NULL
+        );
+        var context = LlvmCompilerContext.EMPTY
+            .addLocal("a", LlvmIr.operandInt(1))
+            .addLocal("b", LlvmIr.operandInt(2));
+
+        var result = LlvmCompiler.compileInstruction(instruction, context);
+
+        var topLevelEntities = new StringBuilder();
+        forEachInterspersed(
+            result.topLevelEntities(),
+            topLevelEntity -> LlvmIrSerializer.serializeTopLevelEntity(topLevelEntity, topLevelEntities),
+            () -> topLevelEntities.append("\n")
+        );
+        assertThat(topLevelEntities.toString(), equalTo("""
+            define i64 @shed__fun__f__0([2 x i64] * %shed__env__0) {
+                ret i64 0
+            }
+            """));
+
+        var instructions = new StringBuilder();
+        forEachInterspersed(
+            result.instructions(),
+            x -> LlvmIrSerializer.serializeInstruction(x, instructions),
+            () -> instructions.append("\n")
+        );
+        assertThat(instructions.toString(), equalTo("""
+            %size_ptr__0 = getelementptr {i64 ([2 x i64] *) *, [2 x i64]}, ptr null, i64 0
+            %size__0 = ptrtoint {i64 ([2 x i64] *) *, [2 x i64]} * %size_ptr__0 to i64
+            %closurePointer__0 = call {i64 ([2 x i64] *) *, [2 x i64]} * @shed__runtime__alloc(i64 %size__0)
+            %closureFunctionPointerPointer__0 = getelementptr {i64 ([2 x i64] *) *, [2 x i64]}, ptr %closurePointer__0, i64 0, i32 0
+            store i64 ([2 x i64] *) * @shed__fun__f__0, i64 ([2 x i64] *) * * %closureFunctionPointerPointer__0
+            %closureEnvironmentPointer__0 = getelementptr {i64 ([2 x i64] *) *, [2 x i64]}, ptr %closurePointer__0, i64 0, i32 1, i64 0
+            store i64 1, i64 * %closureEnvironmentPointer__0
+            %closureEnvironmentPointer__1 = getelementptr {i64 ([2 x i64] *) *, [2 x i64]}, ptr %closurePointer__0, i64 0, i32 1, i64 1
+            store i64 2, i64 * %closureEnvironmentPointer__1"""));
+
+        assertThat(result.stack(), isSequence(
+            isLlvmVariableLocal("closurePointer__0")
+        ));
+    }
+
     private static Matcher<Map.Entry<String, LlvmOperand>> isLocal(String name, Matcher<LlvmOperand> operand) {
         return allOf(
             has("name", x -> x.getKey(), equalTo(name)),