Extract function for creating a closure

2 files changed, 122 insertions, 107 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 2afcac4a..d0f3424f 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -88,13 +88,6 @@ public class LlvmCompiler {
                 );
             }
             case DefineFunctionInstruction defineFunction -> {
-                var closureTypes = LlvmCompilerClosures.types(defineFunction.freeVariables());
-
-                var topLevelEntityName = context.generateTopLevelEntityName(
-                    LlvmCompilerTopLevelEntitySpace.FUNCTION,
-                    defineFunction.name()
-                );
-
                 var functionContext = context.enterFunction();
                 for (var bodyInstruction : defineFunction.body()) {
                     functionContext = compileInstruction(bodyInstruction, functionContext);
@@ -103,104 +96,12 @@ public class LlvmCompiler {
                 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"
-                        )
-                    ),
-                    body
-                ));
-
-                var closurePointerResult = context.generateVariableLocal(
-                    LlvmCompilerLocalVariableSpace.CLOSURE_POINTER
-                );
-                var closurePointer = closurePointerResult.value();
-                context = closurePointerResult.context();
-
-                var closureSizeResult = LlvmSizeOf.sizeOf(
-                    closureTypes.closureStruct(),
+                yield LlvmCompilerClosures.create(
+                    defineFunction.name(),
+                    defineFunction.freeVariables(),
+                    body,
                     context
                 );
-                var closureSize = closureSizeResult.value();
-                context = closureSizeResult.context();
-
-                context = context.addInstruction(new LlvmCall(
-                    closurePointer,
-                    closureTypes.closurePointer(),
-                    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,
-                    closureTypes.closurePointer(),
-                    closurePointer,
-                    List.of(LlvmTypedOperand.arrayIndex(0), LlvmTypedOperand.structureIndex(0))
-                ));
-
-                context = context.addInstruction(new LlvmStore(
-                    closureTypes.closureFunctionPointer(),
-                    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,
-                        closureTypes.closurePointer(),
-                        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++;
-                }
-
-
-                var closureValueResult = context.generateVariableLocal(
-                    LlvmCompilerLocalVariableSpace.CLOSURE_VALUE
-                );
-                var closureValue = closureValueResult.value();
-                context = closureValueResult.context();
-
-                context = context.addInstruction(new LlvmPtrToInt(
-                    closureValue,
-                    closurePointer,
-                    LlvmCompilerTypes.VALUE_TYPE
-                ));
-
-                context = context.push(closureValue);
-
-                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 5c23da3f..eece1c39 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
@@ -1,11 +1,12 @@
 package org.shedlang.compiler.backends.llvm;
 
-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.LlvmIntToPtr;
-import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmLoad;
+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.*;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmTypedOperand;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableGlobal;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.*;
 
@@ -60,6 +61,119 @@ public class LlvmCompilerClosures {
         return new Types(freeVariables);
     }
 
+    public static LlvmCompilerContext create(
+        String name,
+        Set<String> freeVariables,
+        List<LlvmBasicBlock> body,
+        LlvmCompilerContext context
+    ) {
+        var closureTypes = LlvmCompilerClosures.types(freeVariables);
+
+        var topLevelEntityName = context.generateTopLevelEntityName(
+            LlvmCompilerTopLevelEntitySpace.FUNCTION,
+            name
+        );
+
+        context = context.addTopLevelEntity(new LlvmFunctionDefinition(
+            topLevelEntityName,
+            LlvmCompilerTypes.VALUE_TYPE,
+            List.of(
+                new LlvmParam(
+                    closureTypes.closureEnvironmentPointer(),
+                    // TODO: allocate name properly
+                    "shed__env__0"
+                )
+            ),
+            body
+        ));
+
+        var closurePointerResult = context.generateVariableLocal(
+            LlvmCompilerLocalVariableSpace.CLOSURE_POINTER
+        );
+        var closurePointer = closurePointerResult.value();
+        context = closurePointerResult.context();
+
+        var closureSizeResult = LlvmSizeOf.sizeOf(
+            closureTypes.closureStruct(),
+            context
+        );
+        var closureSize = closureSizeResult.value();
+        context = closureSizeResult.context();
+
+        context = context.addInstruction(new LlvmCall(
+            closurePointer,
+            closureTypes.closurePointer(),
+            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,
+            closureTypes.closurePointer(),
+            closurePointer,
+            List.of(LlvmTypedOperand.arrayIndex(0), LlvmTypedOperand.structureIndex(0))
+        ));
+
+        context = context.addInstruction(new LlvmStore(
+            closureTypes.closureFunctionPointer(),
+            new LlvmVariableGlobal(topLevelEntityName),
+            closureFunctionPointerPointer
+        ));
+
+        var freeVariableIndex = 0;
+        // TODO: change type of free variables to be list?
+        for (var freeVariable : freeVariables) {
+            var closureEnvironmentPointerResult = context.generateVariableLocal(
+                LlvmCompilerLocalVariableSpace.CLOSURE_ENVIRONMENT_POINTER
+            );
+            var closureEnvironmentPointer = closureEnvironmentPointerResult.value();
+            context = closureEnvironmentPointerResult.context();
+
+            context = context.addInstruction(new LlvmGetElementPtr(
+                closureEnvironmentPointer,
+                closureTypes.closurePointer(),
+                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++;
+        }
+
+
+        var closureValueResult = context.generateVariableLocal(
+            LlvmCompilerLocalVariableSpace.CLOSURE_VALUE
+        );
+        var closureValue = closureValueResult.value();
+        context = closureValueResult.context();
+
+        context = context.addInstruction(new LlvmPtrToInt(
+            closureValue,
+            closurePointer,
+            LlvmCompilerTypes.VALUE_TYPE
+        ));
+
+        context = context.push(closureValue);
+
+        return context;
+    }
+
     public static LlvmCompilerContext call(
         LlvmOperand closure,
         LlvmCompilerContext context,