Store free variables into environment

2 files changed, 66 insertions, 10 deletions

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)),