Replace enter/exit function with createFunction()

3 files changed, 66 insertions, 73 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
index ad5ec2d5..2c18ad72 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -20,38 +20,31 @@ public class LlvmCodeGen {
     public static LlvmModule generateMain(NamespaceName main, TypeInfo typeInfo) {
         var context = LlvmCodeGenContext.fromTypeInfo(typeInfo);
 
-        var mainContext = context.enterFunction();
-
-        var namespaceResult = LlvmCodeGenNamespaces.importValue(main, mainContext);
-        var namespace = namespaceResult.value();
-        mainContext = namespaceResult.context();
-
-        var fieldResult = LlvmCodeGenNamespaces.fieldLoad(
-            main,
-            namespace,
-            // TODO: check there is a main function
-            "main",
-            mainContext
-        );
-        var field = fieldResult.value();
-        mainContext = fieldResult.context();
-
-        mainContext = LlvmCodeGenClosures.call(field, List.of(), List.of(), mainContext);
-
-        mainContext = mainContext.addTerminatorInstruction(
-            new LlvmReturn(LlvmTypes.INT_32, new LlvmOperandInt(0))
-        );
-
-        var mainBasicBlocksResult = context.exitFunction(mainContext);
-        var mainBasicBlocks = mainBasicBlocksResult.value();
-        context = mainBasicBlocksResult.context();
-
-        context = context.addTopLevelEntity(new LlvmFunctionDefinition(
-            "main",
-            LlvmTypes.INT_32,
-            List.of(),
-            mainBasicBlocks
-        ));
+        context = context.createFunction(mainContext -> {
+            var namespaceResult = LlvmCodeGenNamespaces.importValue(main, mainContext);
+            var namespace = namespaceResult.value();
+            mainContext = namespaceResult.context();
+
+            var fieldResult = LlvmCodeGenNamespaces.fieldLoad(
+                main,
+                namespace,
+                // TODO: check there is a main function
+                "main",
+                mainContext
+            );
+            var field = fieldResult.value();
+            mainContext = fieldResult.context();
+
+            mainContext = LlvmCodeGenClosures.call(field, List.of(), List.of(), mainContext);
+
+            mainContext = mainContext.addTerminatorInstruction(
+                new LlvmReturn(LlvmTypes.INT_32, new LlvmOperandInt(0))
+            );
+
+            var function = mainContext.toFunction("main", LlvmTypes.INT_32, List.of());
+
+            return mainContext.toResult(function);
+        });
 
         return context.toModule();
     }
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java
index 62212a34..81aa781b 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.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.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;
@@ -11,7 +10,10 @@ 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.*;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.function.Function;
 
 import static org.shedlang.compiler.util.Iterables.sortedBy;
@@ -75,38 +77,36 @@ public class LlvmCodeGenClosures {
             name
         );
 
-        var functionContext = context.enterFunction();
-        var functionClosurePointerResult = functionContext.generateVariableLocal(
-            LlvmCodeGenLocalVariableSpace.CLOSURE_POINTER
-        );
-        var functionClosurePointer = functionClosurePointerResult.value();
-        functionContext = functionClosurePointerResult.context();
-
-        for (var freeVariable : freeVariables) {
-            functionContext = loadClosureFreeVariable(
-                functionClosurePointer,
-                freeVariable,
-                closureTypes,
-                functionContext
+        context = context.createFunction(bodyContext -> {
+            var functionClosurePointerResult = bodyContext.generateVariableLocal(
+                LlvmCodeGenLocalVariableSpace.CLOSURE_POINTER
             );
-        }
-
-        functionContext = compileBody.apply(functionContext);
-        var bodyResult = context.exitFunction(functionContext);
-        var body = bodyResult.value();
-        context = bodyResult.context();
-
-        context = context.addTopLevelEntity(new LlvmFunctionDefinition(
-            topLevelEntityName,
-            LlvmCodeGenTypes.VALUE_TYPE,
-            List.of(
-                new LlvmParam(
-                    closureTypes.closurePointer(),
-                    functionClosurePointer.name()
+            var functionClosurePointer = functionClosurePointerResult.value();
+            bodyContext = functionClosurePointerResult.context();
+
+            for (var freeVariable : freeVariables) {
+                bodyContext = loadClosureFreeVariable(
+                    functionClosurePointer,
+                    freeVariable,
+                    closureTypes,
+                    bodyContext
+                );
+            }
+
+            bodyContext = compileBody.apply(bodyContext);
+
+            var function = bodyContext.toFunction(
+                topLevelEntityName,
+                LlvmCodeGenTypes.VALUE_TYPE,
+                List.of(
+                    new LlvmParam(
+                        closureTypes.closurePointer(),
+                        functionClosurePointer.name()
+                    )
                 )
-            ),
-            body
-        ));
+            );
+            return bodyContext.toResult(function);
+        });
 
         var closurePointerResult = context.generateVariableLocal(
             LlvmCodeGenLocalVariableSpace.CLOSURE_POINTER
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
index 3b744a02..e7dddf02 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
@@ -16,6 +16,7 @@ import org.shedlang.compiler.util.POrderedMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 
 public record LlvmCodeGenContext(
     ModuleContext moduleContext,
@@ -328,20 +329,19 @@ public record LlvmCodeGenContext(
         return new LlvmCodeGenResult<>(value, this);
     }
 
-    public LlvmCodeGenContext enterFunction() {
-        return new LlvmCodeGenContext(
+    public LlvmCodeGenContext createFunction(
+        Function<LlvmCodeGenContext, LlvmCodeGenResult<LlvmFunctionDefinition>> create
+    ) {
+        var bodyContext = new LlvmCodeGenContext(
             this.moduleContext,
             FunctionContext.EMPTY
         );
-    }
 
-    public LlvmCodeGenResult<List<LlvmBasicBlock>> exitFunction(
-        LlvmCodeGenContext functionContext
-    ) {
-        // TODO: merge top-level entities
+        var result = create.apply(bodyContext);
+
         return new LlvmCodeGenContext(
-            functionContext.moduleContext,
+            result.context().moduleContext(),
             this.functionContext
-        ).toResult(functionContext.functionContext.basicBlocks);
+        ).addTopLevelEntity(result.value());
     }
 }