Split LLVM code gen context into module and function

3 files changed, 216 insertions, 121 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 5eb418a9..ad5ec2d5 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -53,7 +53,7 @@ public class LlvmCodeGen {
             mainBasicBlocks
         ));
 
-        return new LlvmModule(context.topLevelEntities());
+        return context.toModule();
     }
 
     public static LlvmModule generateNamespace(
@@ -68,9 +68,7 @@ public class LlvmCodeGen {
 
         context = LlvmCodeGenNamespaces.define(namespaceName, context);
 
-        var topLevelEntities = context.topLevelEntities();
-
-        return new LlvmModule(topLevelEntities);
+        return context.toModule();
     }
 
     private static LlvmCodeGenContext defaultContext(TypeInfo typeInfo) {
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 aedb2c55..dd5e5015 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
@@ -6,6 +6,7 @@ import org.pcollections.PStack;
 import org.pcollections.PVector;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmGlobalDefinition;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmModule;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmTopLevelEntity;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmInstruction;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmTerminatorInstruction;
@@ -15,105 +16,171 @@ import org.shedlang.compiler.typechecker.TypeInfo;
 import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 public record LlvmCodeGenContext(
-    TypeInfo typeInfo,
-    PStack<LlvmOperand> stack,
-    PMap<String, LlvmOperand> shedLocalNameToOperand,
-    PVector<LlvmBasicBlock> basicBlocks,
-    PVector<LlvmInstruction> instructions,
-    POrderedMap<String, LlvmTopLevelEntity> topLevelEntitiesByName,
-    NameAllocator<LlvmCodeGenLocalVariableSpace> localVariableNames
+    ModuleContext moduleContext,
+    FunctionContext functionContext
 ) {
-    public static LlvmCodeGenContext EMPTY = fromTypeInfo(TypeInfo.EMPTY);
+    record ModuleContext(
+        TypeInfo typeInfo,
+        POrderedMap<String, LlvmTopLevelEntity> topLevelEntitiesByName
+    ) {
+        public static ModuleContext fromTypeInfo(TypeInfo typeInfo) {
+            return new ModuleContext(typeInfo, POrderedMap.empty());
+        }
 
-    public static LlvmCodeGenContext fromTypeInfo(TypeInfo typeInfo) {
-        return new LlvmCodeGenContext(
-            typeInfo,
+        public ModuleContext addTopLevelEntity(LlvmTopLevelEntity topLevelEntity) {
+            return new ModuleContext(
+                typeInfo,
+                topLevelEntitiesByName.plus(topLevelEntity.name(), topLevelEntity)
+            );
+        }
+
+        public String generateTopLevelEntityName(
+            LlvmCodeGenTopLevelEntitySpace space,
+            Optional<String> name
+        ) {
+            return NameAllocator.generateName(
+                topLevelEntitiesByName::containsKey,
+                "shed__" + space,
+                name
+            );
+        }
+
+        public boolean hasTopLevelEntity(String name) {
+            return topLevelEntitiesByName.containsKey(name);
+        }
+
+        public LlvmModule toModule() {
+            return new LlvmModule(topLevelEntitiesByName.toValues());
+        }
+
+        public <T> Result<T> toResult(T value) {
+            return new Result<>(value, this);
+        }
+
+        record Result<T>(T value, ModuleContext context) {}
+    }
+
+    record FunctionContext(
+        PStack<LlvmOperand> stack,
+        PMap<String, LlvmOperand> shedLocalNameToOperand,
+        PVector<LlvmBasicBlock> basicBlocks,
+        PVector<LlvmInstruction> instructions,
+        NameAllocator<LlvmCodeGenLocalVariableSpace> localVariableNames
+    ) {
+        public static final FunctionContext EMPTY = new FunctionContext(
             Empty.stack(),
             Empty.map(),
             Empty.vector(),
             Empty.vector(),
-            POrderedMap.empty(),
             NameAllocator.empty()
         );
-    }
 
-    public LlvmCodeGenContext push(LlvmOperand operand) {
-        return new LlvmCodeGenContext(
-            typeInfo,
-            stack.plus(operand),
-            shedLocalNameToOperand,
-            basicBlocks,
-            instructions,
-            topLevelEntitiesByName,
-            localVariableNames
-        );
-    }
+        public FunctionContext push(LlvmOperand operand) {
+            return new FunctionContext(
+                stack.plus(operand),
+                shedLocalNameToOperand,
+                basicBlocks,
+                instructions,
+                localVariableNames
+            );
+        }
 
-    public LlvmCodeGenResult<LlvmOperand> pop() {
-        return new LlvmCodeGenContext(
-            typeInfo,
-            stack.minus(0),
-            shedLocalNameToOperand,
-            basicBlocks,
-            instructions,
-            topLevelEntitiesByName,
-            localVariableNames
-        ).toResult(stack.get(0));
-    }
+        public Result<LlvmOperand> pop() {
+            return new FunctionContext(
+                stack.minus(0),
+                shedLocalNameToOperand,
+                basicBlocks,
+                instructions,
+                localVariableNames
+            ).toResult(stack.get(0));
+        }
 
-    public LlvmCodeGenContext addShedLocal(String name, LlvmOperand value) {
-        return new LlvmCodeGenContext(
-            typeInfo,
-            stack,
-            shedLocalNameToOperand.plus(name, value),
-            basicBlocks,
-            instructions,
-            topLevelEntitiesByName,
-            localVariableNames
-        );
-    }
+        public FunctionContext addShedLocal(String name, LlvmOperand value) {
+            return new FunctionContext(
+                stack,
+                shedLocalNameToOperand.plus(name, value),
+                basicBlocks,
+                instructions,
+                localVariableNames
+            );
+        }
 
-    public LlvmOperand shedLocal(String name) {
-        var local = shedLocalNameToOperand.get(name);
-        if (local == null) {
-            // TODO: better error
-            throw new RuntimeException("unrecognised local: " + name);
-        } else {
-            return local;
+        public LlvmOperand shedLocal(String name) {
+            var local = shedLocalNameToOperand.get(name);
+            if (local == null) {
+                // TODO: better error
+                throw new RuntimeException("unrecognised local: " + name);
+            } else {
+                return local;
+            }
+        }
+
+        public FunctionContext addInstruction(LlvmInstruction instruction) {
+            return new FunctionContext(
+                stack,
+                shedLocalNameToOperand,
+                basicBlocks,
+                instructions.plus(instruction),
+                localVariableNames
+            );
+        }
+
+        public FunctionContext addTerminatorInstruction(
+            LlvmTerminatorInstruction terminatorInstruction
+        ) {
+            return new FunctionContext(
+                stack,
+                shedLocalNameToOperand,
+                basicBlocks.plus(new LlvmBasicBlock(
+                    Optional.empty(),
+                    instructions,
+                    terminatorInstruction
+                )),
+                Empty.vector(),
+                localVariableNames
+            );
+        }
+
+        public Result<LlvmVariableLocal> generateVariableLocal(
+            LlvmCodeGenLocalVariableSpace space,
+            Optional<String> name
+        ) {
+            var result = localVariableNames.generateName(space, name);
+            return new FunctionContext(
+                stack,
+                shedLocalNameToOperand,
+                basicBlocks,
+                instructions,
+                result.allocator()
+            ).toResult(new LlvmVariableLocal(result.name()));
         }
+
+        public <T> Result<T> toResult(T value) {
+            return new Result<>(value, this);
+        }
+
+        record Result<T>(T value, FunctionContext context) {}
     }
 
-    public LlvmCodeGenContext addInstruction(LlvmInstruction instruction) {
+    public static LlvmCodeGenContext EMPTY = fromTypeInfo(TypeInfo.EMPTY);
+
+    public static LlvmCodeGenContext fromTypeInfo(TypeInfo typeInfo) {
         return new LlvmCodeGenContext(
-            typeInfo,
-            stack,
-            shedLocalNameToOperand,
-            basicBlocks,
-            instructions.plus(instruction),
-            topLevelEntitiesByName,
-            localVariableNames
+            ModuleContext.fromTypeInfo(typeInfo),
+            FunctionContext.EMPTY
         );
     }
 
-    public LlvmCodeGenContext addTerminatorInstruction(
-        LlvmTerminatorInstruction terminatorInstruction
-    ) {
-        return new LlvmCodeGenContext(
-            typeInfo,
-            stack,
-            shedLocalNameToOperand,
-            basicBlocks.plus(new LlvmBasicBlock(
-                Optional.empty(),
-                instructions,
-                terminatorInstruction
-            )),
-            Empty.vector(),
-            topLevelEntitiesByName,
-            localVariableNames
-        );
+    public TypeInfo typeInfo() {
+        return moduleContext.typeInfo;
+    }
+
+    public List<LlvmTopLevelEntity> topLevelEntities() {
+        return moduleContext.topLevelEntitiesByName.toValues();
     }
 
     public LlvmCodeGenResult<LlvmGlobalDefinition> addGlobalDefinition(
@@ -133,13 +200,8 @@ public record LlvmCodeGenContext(
 
     public LlvmCodeGenContext addTopLevelEntity(LlvmTopLevelEntity topLevelEntity) {
         return new LlvmCodeGenContext(
-            typeInfo,
-            stack,
-            shedLocalNameToOperand,
-            basicBlocks,
-            instructions,
-            topLevelEntitiesByName.plus(topLevelEntity.name(), topLevelEntity),
-            localVariableNames
+            this.moduleContext.addTopLevelEntity(topLevelEntity),
+            this.functionContext
         );
     }
 
@@ -155,19 +217,71 @@ public record LlvmCodeGenContext(
         LlvmCodeGenTopLevelEntitySpace space,
         Optional<String> name
     ) {
-        return NameAllocator.generateName(
-            topLevelEntitiesByName::containsKey,
-            "shed__" + space,
-            name
-        );
+        return moduleContext.generateTopLevelEntityName(space, name);
     }
 
     public boolean hasTopLevelEntity(String name) {
-        return topLevelEntitiesByName.containsKey(name);
+        return moduleContext.hasTopLevelEntity(name);
     }
 
-    public List<LlvmTopLevelEntity> topLevelEntities() {
-        return topLevelEntitiesByName.toValues();
+    public LlvmModule toModule() {
+        return moduleContext.toModule();
+    }
+
+    public PStack<LlvmOperand> stack() {
+        return functionContext.stack;
+    }
+
+    public LlvmCodeGenContext push(LlvmOperand operand) {
+        return new LlvmCodeGenContext(
+            moduleContext,
+            functionContext.push(operand)
+        );
+    }
+
+    public LlvmCodeGenResult<LlvmOperand> pop() {
+        var result = functionContext.pop();
+
+        return new LlvmCodeGenContext(moduleContext, result.context).toResult(result.value);
+    }
+
+    public LlvmCodeGenContext addShedLocal(String name, LlvmOperand value) {
+        return new LlvmCodeGenContext(
+            this.moduleContext,
+            this.functionContext.addShedLocal(name, value)
+        );
+    }
+
+    public LlvmOperand shedLocal(String name) {
+        return this.functionContext.shedLocal(name);
+    }
+
+    public Map<String, LlvmOperand> shedLocalNameToOperand() {
+        return this.functionContext.shedLocalNameToOperand;
+    }
+
+    public List<LlvmBasicBlock> basicBlocks() {
+        return functionContext.basicBlocks;
+    }
+
+    public List<LlvmInstruction> instructions() {
+        return functionContext.instructions;
+    }
+
+    public LlvmCodeGenContext addInstruction(LlvmInstruction instruction) {
+        return new LlvmCodeGenContext(
+            this.moduleContext,
+            this.functionContext.addInstruction(instruction)
+        );
+    }
+
+    public LlvmCodeGenContext addTerminatorInstruction(
+        LlvmTerminatorInstruction terminatorInstruction
+    ) {
+        return new LlvmCodeGenContext(
+            this.moduleContext,
+            this.functionContext.addTerminatorInstruction(terminatorInstruction)
+        );
     }
 
     public LlvmCodeGenResult<LlvmVariableLocal> generateVariableLocal(
@@ -187,16 +301,8 @@ public record LlvmCodeGenContext(
         LlvmCodeGenLocalVariableSpace space,
         Optional<String> name
     ) {
-        var result = localVariableNames.generateName(space, name);
-        return new LlvmCodeGenContext(
-            typeInfo,
-            stack,
-            shedLocalNameToOperand,
-            basicBlocks,
-            instructions,
-            topLevelEntitiesByName,
-            result.allocator()
-        ).toResult(new LlvmVariableLocal(result.name()));
+        var result = this.functionContext.generateVariableLocal(space, name);
+        return new LlvmCodeGenContext(this.moduleContext, result.context).toResult(result.value);
     }
 
     public <T> LlvmCodeGenResult<T> toResult(T value) {
@@ -205,13 +311,8 @@ public record LlvmCodeGenContext(
 
     public LlvmCodeGenContext enterFunction() {
         return new LlvmCodeGenContext(
-            typeInfo,
-            Empty.stack(),
-            Empty.map(),
-            Empty.vector(),
-            Empty.vector(),
-            topLevelEntitiesByName,
-            NameAllocator.empty()
+            this.moduleContext,
+            FunctionContext.EMPTY
         );
     }
 
@@ -220,13 +321,8 @@ public record LlvmCodeGenContext(
     ) {
         // TODO: merge top-level entities
         return new LlvmCodeGenContext(
-            typeInfo,
-            stack,
-            shedLocalNameToOperand,
-            basicBlocks,
-            instructions,
-            functionContext.topLevelEntitiesByName,
-            localVariableNames
-        ).toResult(functionContext.basicBlocks);
+            functionContext.moduleContext,
+            this.functionContext
+        ).toResult(functionContext.functionContext.basicBlocks);
     }
 }
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 1c60bab9..c595838d 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenNamespaces.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenNamespaces.java
@@ -103,6 +103,7 @@ public class LlvmCodeGenNamespaces {
             new LlvmVariableGlobal(globalNameNamespaceIsInited(namespaceName))
         ));
         context = context.addTerminatorInstruction(new LlvmReturnVoid());
+        // TODO: add toFunction()
         var basicBlocks = context.basicBlocks();
         var namespaceInitFunction = new LlvmFunctionDefinition(
             globalNameNamespaceInit(namespaceName),