Expand function definition

2 files changed, 60 insertions, 9 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 4ee42fa8..694f3ff9 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -2,9 +2,12 @@ package org.shedlang.compiler.backends.llvm;
 
 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.LlvmCall;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmGetElementPtr;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmReturn;
-import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
-import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperandInt;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmStore;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.*;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.types.NamespaceName;
 
@@ -66,7 +69,7 @@ public class LlvmCompiler {
                 var body = bodyResult.value();
                 context = bodyResult.context();
 
-                yield context.addTopLevelEntity(new LlvmFunctionDefinition(
+                context = context.addTopLevelEntity(new LlvmFunctionDefinition(
                     topLevelEntityName,
                     LlvmCompilerTypes.VALUE_TYPE,
                     List.of(
@@ -78,6 +81,52 @@ public class LlvmCompiler {
                     ),
                     body
                 ));
+
+                var closurePointerResult = context.generateVariableLocal(
+                    LlvmCompilerLocalVariableSpace.CLOSURE_POINTER
+                );
+                var closurePointer = closurePointerResult.value();
+                context = closurePointerResult.context();
+
+                var closureSizeResult = LlvmSizeOf.sizeOf(
+                    LlvmCompilerClosures.typeClosureStruct(Optional.of(defineFunction.freeVariables())),
+                    context
+                );
+                var closureSize = closureSizeResult.value();
+                context = closureSizeResult.context();
+
+                context = context.addInstruction(new LlvmCall(
+                    closurePointer,
+                    LlvmCompilerClosures.typeClosurePointer(),
+                    new LlvmVariableGlobal("shed__runtime__alloc"),
+                    List.of(
+                        new LlvmTypedOperand(LlvmTypes.INT_64, closureSize)
+                    )
+                ));
+
+                var closureFunctionPointerPointerResult = context.generateVariableLocal(
+                    LlvmCompilerLocalVariableSpace.CLOSURE_FUNCTION_POINTER_POINTER
+                );
+                var closureFunctionPointerPointer = closureFunctionPointerPointerResult.value();
+                context = closureFunctionPointerPointerResult.context();
+
+                context = context.addInstruction(new LlvmGetElementPtr(
+                    closureFunctionPointerPointer,
+                    LlvmCompilerClosures.typeClosurePointer(),
+                    closurePointer,
+                    List.of(LlvmTypedOperand.arrayIndex(0), LlvmTypedOperand.structureIndex(0))
+                ));
+
+                context = context.addInstruction(new LlvmStore(
+                    LlvmCompilerClosures.typeClosureFunctionPointer(),
+                    new LlvmVariableGlobal(topLevelEntityName),
+                    closureFunctionPointerPointer
+                ));
+
+                // TODO: ptrtoint
+                context = context.push(closurePointer);
+
+                yield context;
             }
             case DiscardInstruction discardInstruction ->
                 context.pop().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 a579e6a7..9ea57ec9 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
@@ -11,19 +11,21 @@ import org.shedlang.compiler.backends.llvm.ir.ast.types.*;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
+import java.util.Set;
 
 public class LlvmCompilerClosures {
     private LlvmCompilerClosures() {
     }
 
     public static LlvmTypePointer typeClosurePointer() {
-        return LlvmTypes.pointer(typeClosureStruct());
+        return LlvmTypes.pointer(typeClosureStruct(Optional.empty()));
     }
 
-    public static LlvmTypeStructure typeClosureStruct() {
+    public static LlvmTypeStructure typeClosureStruct(Optional<Set<String>> freeVariables) {
         return LlvmTypes.structure(List.of(
             typeClosureFunctionPointer(),
-            typeClosureEnvironment()
+            typeClosureEnvironment(freeVariables)
         ));
     }
 
@@ -39,11 +41,11 @@ public class LlvmCompilerClosures {
     }
 
     public static LlvmTypePointer typeClosureEnvironmentPointer() {
-        return LlvmTypes.pointer(typeClosureEnvironment());
+        return LlvmTypes.pointer(typeClosureEnvironment(Optional.empty()));
     }
 
-    public static LlvmTypeArray typeClosureEnvironment() {
-        return LlvmTypes.array(0, LlvmCompilerTypes.VALUE_TYPE);
+    public static LlvmTypeArray typeClosureEnvironment(Optional<Set<String>> freeVariables) {
+        return LlvmTypes.array(freeVariables.map(x -> x.size()).orElse(0), LlvmCompilerTypes.VALUE_TYPE);
     }
 
     public static LlvmCompilerContext call(