Pass function to create init body into LlvmCodeGenNamespaces.define()

2 files changed, 45 insertions, 26 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 2c18ad72..50c05a9c 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -55,11 +55,18 @@ public class LlvmCodeGen {
         TypeInfo typeInfo
     ) {
         var context = defaultContext(typeInfo);
-        for (var instruction : instructions) {
-            context = generateInstruction(instruction, context);
-        }
 
-        context = LlvmCodeGenNamespaces.define(namespaceName, context);
+        context = LlvmCodeGenNamespaces.define(
+            namespaceName,
+            initContext -> {
+                for (var instruction : instructions) {
+                    initContext = generateInstruction(instruction, initContext);
+                }
+
+                return initContext;
+            },
+            context
+        );
 
         return context.toModule();
     }
@@ -78,10 +85,6 @@ public class LlvmCodeGen {
                 false,
                 false,
                 false
-            ))
-            .addShedLocal("print", new LlvmOperandPtrToInt(
-                new LlvmVariableGlobal("shed__runtime__print_closure"),
-                LlvmCodeGenTypes.VALUE_TYPE
             ));
     }
 
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenNamespaces.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenNamespaces.java
index 726a77f7..62ecaa9f 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenNamespaces.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenNamespaces.java
@@ -16,6 +16,7 @@ import org.shedlang.compiler.types.NamespaceTypeInfo;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.stream.IntStream;
 
 public class LlvmCodeGenNamespaces {
@@ -95,19 +96,41 @@ public class LlvmCodeGenNamespaces {
 
     public static LlvmCodeGenContext define(
         NamespaceName namespaceName,
+        Function<LlvmCodeGenContext, LlvmCodeGenContext> generateInitBody,
         LlvmCodeGenContext context
     ) {
-        context = context.addInstruction(new LlvmStore(
-            LlvmTypes.INT_1,
-            new LlvmOperandInt(1),
-            new LlvmVariableGlobal(globalNameNamespaceIsInited(namespaceName))
+        context = context.addTopLevelEntity(defineNamespaceIsInited(namespaceName));
+
+        context = context.addTopLevelEntity(defineNamespaceValue(
+            namespaceName,
+            context.typeInfo().namespaceTypeInfo(namespaceName)
         ));
-        context = context.addTerminatorInstruction(new LlvmReturnVoid());
-        var namespaceInitFunction = context.toFunction(
-            globalNameNamespaceInit(namespaceName),
-            LlvmTypes.VOID,
-            List.of()
-        );
+
+        context = context.createFunction(initContext -> {
+            // TODO: handle default locals more gracefully
+            initContext = initContext.addShedLocal("print", new LlvmOperandPtrToInt(
+                new LlvmVariableGlobal("shed__runtime__print_closure"),
+                LlvmCodeGenTypes.VALUE_TYPE
+            ));
+
+            initContext = generateInitBody.apply(initContext);
+
+            initContext = initContext.addInstruction(new LlvmStore(
+                LlvmTypes.INT_1,
+                new LlvmOperandInt(1),
+                new LlvmVariableGlobal(globalNameNamespaceIsInited(namespaceName))
+            ));
+
+            initContext = initContext.addTerminatorInstruction(new LlvmReturnVoid());
+
+            var namespaceInitFunction = initContext.toFunction(
+                globalNameNamespaceInit(namespaceName),
+                LlvmTypes.VOID,
+                List.of()
+            );
+
+            return initContext.toResult(namespaceInitFunction);
+        });
 
         var namespaceImportFunction = new LlvmFunctionDefinition(
             globalNameNamespaceImport(namespaceName),
@@ -152,14 +175,7 @@ public class LlvmCodeGenNamespaces {
             )
         );
 
-        return context
-            .addTopLevelEntity(defineNamespaceIsInited(namespaceName))
-            .addTopLevelEntity(defineNamespaceValue(
-                namespaceName,
-                context.typeInfo().namespaceTypeInfo(namespaceName)
-            ))
-            .addTopLevelEntity(namespaceInitFunction)
-            .addTopLevelEntity(namespaceImportFunction);
+        return context.addTopLevelEntity(namespaceImportFunction);
     }
 
     private static LlvmGlobalDefinition defineNamespaceIsInited(