Load free variables from closure

3 files changed, 64 insertions, 21 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 d0f3424f..a087b75e 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -88,18 +88,15 @@ public class LlvmCompiler {
                 );
             }
             case DefineFunctionInstruction defineFunction -> {
-                var functionContext = context.enterFunction();
-                for (var bodyInstruction : defineFunction.body()) {
-                    functionContext = compileInstruction(bodyInstruction, functionContext);
-                }
-                var bodyResult = context.exitFunction(functionContext);
-                var body = bodyResult.value();
-                context = bodyResult.context();
-
                 yield LlvmCompilerClosures.create(
                     defineFunction.name(),
                     defineFunction.freeVariables(),
-                    body,
+                    functionContext -> {
+                        for (var bodyInstruction : defineFunction.body()) {
+                            functionContext = compileInstruction(bodyInstruction, functionContext);
+                        }
+                        return functionContext;
+                    },
                     context
                 );
             }
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 b88a8401..22c55ce2 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
@@ -2,7 +2,6 @@ package org.shedlang.compiler.backends.llvm;
 
 import org.shedlang.compiler.ast.typed.FreeVariable;
 import org.shedlang.compiler.ast.typed.FreeVariables;
-import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmFunctionDefinition;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmParam;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.*;
@@ -15,6 +14,7 @@ import org.shedlang.compiler.backends.llvm.ir.ast.types.*;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 
 public class LlvmCompilerClosures {
     private LlvmCompilerClosures() {
@@ -65,7 +65,7 @@ public class LlvmCompilerClosures {
     public static LlvmCompilerContext create(
         String name,
         FreeVariables freeVariables,
-        List<LlvmBasicBlock> body,
+        Function<LlvmCompilerContext, LlvmCompilerContext> compileBody,
         LlvmCompilerContext context
     ) {
         var closureTypes = LlvmCompilerClosures.types(freeVariables);
@@ -75,14 +75,34 @@ public class LlvmCompilerClosures {
             name
         );
 
+        var functionContext = context.enterFunction();
+        var functionClosurePointerResult = context.generateVariableLocal(
+            LlvmCompilerLocalVariableSpace.CLOSURE_POINTER
+        );
+        var functionClosurePointer = functionClosurePointerResult.value();
+        functionContext = functionClosurePointerResult.context();
+
+        for (var freeVariable : freeVariables) {
+            functionContext = loadClosureFreeVariable(
+                functionClosurePointer,
+                freeVariable,
+                closureTypes,
+                functionContext
+            );
+        }
+
+        functionContext = compileBody.apply(functionContext);
+        var bodyResult = context.exitFunction(functionContext);
+        var body = bodyResult.value();
+        context = bodyResult.context();
+
         context = context.addTopLevelEntity(new LlvmFunctionDefinition(
             topLevelEntityName,
             LlvmCompilerTypes.VALUE_TYPE,
             List.of(
                 new LlvmParam(
-                    closureTypes.closureEnvironmentPointer(),
-                    // TODO: allocate name properly
-                    "shed__env__0"
+                    closureTypes.closurePointer(),
+                    functionClosurePointer.name()
                 )
             ),
             body
@@ -116,7 +136,7 @@ public class LlvmCompilerClosures {
         );
 
         for (var freeVariable : freeVariables) {
-            context = storeClosureFreeVariablePointer(
+            context = storeClosureFreeVariable(
                 closurePointer,
                 freeVariable,
                 closureTypes,
@@ -146,10 +166,6 @@ public class LlvmCompilerClosures {
         var functionPointer = functionPointerResult.value();
         context = functionPointerResult.context();
 
-        var environmentPointerResult = getClosureEnvironmentPointer(closurePointer, context);
-        var environmentPointer = environmentPointerResult.value();
-        context = environmentPointerResult.context();
-
         var callTargetResult = context.generateVariableLocal(
             LlvmCompilerLocalVariableSpace.TEMPORARY
         );
@@ -157,7 +173,7 @@ public class LlvmCompilerClosures {
         context = callTargetResult.context();
 
         var fullArgs = new ArrayList<LlvmTypedOperand>();
-        fullArgs.add(new LlvmTypedOperand(CALLABLE_TYPES.closureEnvironmentPointer(), environmentPointer));
+        fullArgs.add(new LlvmTypedOperand(CALLABLE_TYPES.closurePointer(), closurePointer));
         for (var arg : args) {
             fullArgs.add(new LlvmTypedOperand(LlvmCompilerTypes.VALUE_TYPE, arg));
         }
@@ -295,7 +311,7 @@ public class LlvmCompilerClosures {
         return context.toResult(environmentPointer);
     }
 
-    private static LlvmCompilerContext storeClosureFreeVariablePointer(
+    private static LlvmCompilerContext storeClosureFreeVariable(
         LlvmOperand closurePointer,
         FreeVariable freeVariable,
         Types closureTypes,
@@ -317,6 +333,35 @@ public class LlvmCompilerClosures {
         ));
     }
 
+    private static LlvmCompilerContext loadClosureFreeVariable(
+        LlvmOperand closurePointer,
+        FreeVariable freeVariable,
+        Types closureTypes,
+        LlvmCompilerContext context
+    ) {
+        var closureFreeVariablePointerResult = getClosureFreeVariablePointer(
+            closurePointer,
+            freeVariable,
+            closureTypes,
+            context
+        );
+        var closureFreeVariablePointer = closureFreeVariablePointerResult.value();
+        context = closureFreeVariablePointerResult.context();
+
+        var closureFreeVariableResult = context.generateVariableLocal(
+            LlvmCompilerLocalVariableSpace.CLOSURE_FREE_VARIABLE,
+            freeVariable.name()
+        );
+        var closureFreeVariable = closureFreeVariableResult.value();
+        context = closureFreeVariableResult.context();
+
+        return context.addInstruction(new LlvmLoad(
+            closureFreeVariable,
+            LlvmCompilerTypes.VALUE_TYPE,
+            closureFreeVariablePointer
+        ));
+    }
+
     private static LlvmCompilerResult<LlvmVariableLocal> getClosureFreeVariablePointer(
         LlvmOperand closurePointer,
         FreeVariable freeVariable,
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerLocalVariableSpace.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerLocalVariableSpace.java
index 511e24f4..fd7736e3 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerLocalVariableSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerLocalVariableSpace.java
@@ -7,6 +7,7 @@ public enum LlvmCompilerLocalVariableSpace {
     CLOSURE_FUNCTION_POINTER_POINTER("closureFunctionPointerPointer"),
     CLOSURE_ENVIRONMENT_POINTER("closureEnvironmentPointer"),
     CLOSURE_FREE_VARIABLE_POINTER("closureFreeVariablePointer"),
+    CLOSURE_FREE_VARIABLE("closureFreeVariable"),
     FIELD_POINTER("fieldPointer"),
     SIZE_PTR("sizePtr"),
     SIZE("size"),