Add global context for LLVM code gen

2 files changed, 103 insertions, 58 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
index b40c6455..96cd7e87 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
@@ -49,13 +49,20 @@ public class LlvmBackend implements Backend {
         TargetDir targetDir
     ) throws IOException, InterruptedException {
         var objectFilePaths = new ArrayList<Path>();
+        var globalContext = LlvmPatioCodeGen.GlobalContext.EMPTY;
+
         for (var namespace : namespaces) {
             var patioCompilationUnit = PatioCodeGen.generateNamespace(
                 namespace,
                 PatioCodeGenContext.forTypeInfos(typeInfos)
             );
 
-            var llvmModule = LlvmPatioCodeGen.generateCompilationUnit(patioCompilationUnit);
+            var llvmModuleResult = LlvmPatioCodeGen.generateCompilationUnit(
+                patioCompilationUnit,
+                globalContext
+            );
+            var llvmModule = llvmModuleResult.value();
+            globalContext = llvmModuleResult.context();
 
             var objectPath = compileLlvmModule(
                 llvmModule,
@@ -68,7 +75,7 @@ public class LlvmBackend implements Backend {
         for (var bin : config.bins()) {
             var linkInputs = new ArrayList<>(objectFilePaths);
 
-            var mainModule = this.generateMain(bin.mainNamespace(), typeInfos);
+            var mainModule = this.generateMain(bin.mainNamespace(), typeInfos, globalContext);
             var mainObjectPath = compileLlvmModule(mainModule, targetDir.bin(bin.name()));
             linkInputs.add(mainObjectPath);
 
@@ -92,7 +99,11 @@ public class LlvmBackend implements Backend {
         runInheritIo(gccArgs);
     }
 
-    private LlvmModule generateMain(NamespaceName main, TypeInfos typeInfos) {
+    private LlvmModule generateMain(
+        NamespaceName main,
+        TypeInfos typeInfos,
+        LlvmPatioCodeGen.GlobalContext globalContext
+    ) {
         // TODO: use more appropriate source
         var source = Sources.NULL;
 
@@ -104,7 +115,11 @@ public class LlvmBackend implements Backend {
 
         var patioCompilationUnit = context.toCompilationUnit(source);
 
-        var llvmModule = LlvmPatioCodeGen.generateCompilationUnit(patioCompilationUnit);
+        var llvmModuleResult = LlvmPatioCodeGen.generateCompilationUnit(
+            patioCompilationUnit,
+            globalContext
+        );
+        var llvmModule = llvmModuleResult.value();
 
         llvmModule = llvmModule.addTopLevelEntity(new LlvmFunctionDefinition(
             "main",
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
index 097eb2d8..32ea406a 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -78,7 +78,10 @@ public class LlvmPatioCodeGen {
         )
     );
 
-    public static LlvmModule generateCompilationUnit(PatioCompilationUnit patioCompilationUnit) {
+    public static GlobalContext.Result<LlvmModule> generateCompilationUnit(
+        PatioCompilationUnit patioCompilationUnit,
+        GlobalContext globalContext
+    ) {
         var context = ModuleContext.EMPTY
             .addTopLevelEntity(ALLOC_DECLARATION)
             .addTopLevelEntity(SETJMP_DECLARATION)
@@ -118,7 +121,9 @@ public class LlvmPatioCodeGen {
             context = generateCompilationUnitElement(patioCompilationUnitElement, context);
         }
 
-        return context.toModule();
+        var module = context.toModule();
+
+        return context.globalContext.toResult(module);
     }
 
     private static ModuleContext generateCompilationUnitElement(
@@ -1207,58 +1212,99 @@ public class LlvmPatioCodeGen {
             .toResult(size);
     }
 
+    public static class GlobalContext {
+        public static final GlobalContext EMPTY = new GlobalContext(P.map());
+
+        private final PMap<String, Integer> exceptionTypeNameToTag;
+
+        private GlobalContext(
+            PMap<String, Integer> exceptionTypeNameToTag
+        ) {
+            this.exceptionTypeNameToTag = exceptionTypeNameToTag;
+        }
+
+        public GlobalContext allocateExceptionTag(String exceptionTypeName) {
+            var tag = this.exceptionTypeNameToTag.size() + 1;
+            return new GlobalContext(
+                this.exceptionTypeNameToTag.plus(exceptionTypeName, tag)
+            );
+        }
+
+        public int exceptionTag(String exceptionTypeName) {
+            return this.exceptionTypeNameToTag.get(exceptionTypeName);
+        }
+
+        private ModuleContext enterModule() {
+            return ModuleContext.of(this);
+        }
+
+        private <T> Result<T> toResult(T value) {
+            return new Result<>(value, this);
+        }
+
+        public record Result<T>(T value, GlobalContext context) {
+        }
+    }
+
     private static class ModuleContext {
-        private static final ModuleContext EMPTY = new ModuleContext(
-            P.map(),
-            P.vector(),
-            P.map()
-        );
+        private static final ModuleContext EMPTY = of(GlobalContext.EMPTY);
 
+        public static ModuleContext of(GlobalContext globalContext) {
+            return new ModuleContext(globalContext, P.map(), P.vector());
+        }
+
+        private final GlobalContext globalContext;
         private final PMap<String, LlvmType> functionNameToReturnType;
         private final PVector<LlvmTopLevelEntity> topLevelEntities;
-        private final PMap<String, Integer> exceptionTypeNameToTag;
 
         private ModuleContext(
+            GlobalContext globalContext,
             PMap<String, LlvmType> functionNameToReturnType,
-            PVector<LlvmTopLevelEntity> topLevelEntities,
-            PMap<String, Integer> exceptionTypeNameToTag
+            PVector<LlvmTopLevelEntity> topLevelEntities
         ) {
+            this.globalContext = globalContext;
             this.functionNameToReturnType = functionNameToReturnType;
             this.topLevelEntities = topLevelEntities;
-            this.exceptionTypeNameToTag = exceptionTypeNameToTag;
+        }
+
+        private LlvmModule toModule() {
+            return new LlvmModule(topLevelEntities);
         }
 
         private ModuleContext addFunctionNameToReturnType(String functionName, LlvmType returnType) {
             return new ModuleContext(
+                this.globalContext,
                 this.functionNameToReturnType.plus(functionName, returnType),
-                this.topLevelEntities,
-                this.exceptionTypeNameToTag
+                this.topLevelEntities
             );
         }
 
-        private LlvmModule toModule() {
-            return new LlvmModule(topLevelEntities);
+        private LlvmType functionNameToReturnType(String functionName) {
+            return this.functionNameToReturnType.get(functionName);
         }
 
         public ModuleContext addTopLevelEntity(LlvmTopLevelEntity topLevelEntity) {
             return new ModuleContext(
+                this.globalContext,
                 this.functionNameToReturnType,
-                this.topLevelEntities.plus(topLevelEntity),
-                this.exceptionTypeNameToTag
+                this.topLevelEntities.plus(topLevelEntity)
             );
         }
 
         public ModuleContext allocateExceptionTag(String exceptionTypeName) {
-            var tag = this.exceptionTypeNameToTag.size() + 1;
             return new ModuleContext(
+                this.globalContext.allocateExceptionTag(exceptionTypeName),
                 this.functionNameToReturnType,
-                this.topLevelEntities,
-                this.exceptionTypeNameToTag.plus(exceptionTypeName, tag)
+                this.topLevelEntities
             );
         }
 
+        public int exceptionTag(String exceptionTypeName) {
+            return this.globalContext.exceptionTag(exceptionTypeName);
+        }
+
         public FunctionContext enterFunction() {
-            return FunctionContext.of(functionNameToReturnType, exceptionTypeNameToTag);
+            return FunctionContext.of(this);
         }
 
         private <T> Result<T> toResult(T value) {
@@ -1270,14 +1316,9 @@ public class LlvmPatioCodeGen {
     }
 
     private static class FunctionContext {
-        // TODO: just accept module context?
-        private static FunctionContext of(
-            PMap<String, LlvmType> functionNameToType,
-            PMap<String, Integer> exceptionTypeNameToTag
-        ) {
+        private static FunctionContext of(ModuleContext moduleContext) {
             return new FunctionContext(
-                functionNameToType,
-                exceptionTypeNameToTag,
+                moduleContext,
                 P.vector(),
                 Optional.empty(),
                 P.vector(),
@@ -1288,8 +1329,7 @@ public class LlvmPatioCodeGen {
             );
         }
 
-        private final PMap<String, LlvmType> functionNameToReturnType;
-        private final PMap<String, Integer> exceptionTypeNameToTag;
+        private final ModuleContext moduleContext;
         private final PVector<LlvmBasicBlock> basicBlocks;
         private final Optional<String> basicBlockLabel;
         private final PVector<LlvmInstruction> instructions;
@@ -1299,8 +1339,8 @@ public class LlvmPatioCodeGen {
         private final PStack<LoopLabels> loopLabels;
 
         private FunctionContext(
-            PMap<String, LlvmType> functionNameToReturnType,
-            PMap<String, Integer> exceptionTypeNameToTag, PVector<LlvmBasicBlock> basicBlocks,
+            ModuleContext moduleContext,
+            PVector<LlvmBasicBlock> basicBlocks,
             Optional<String> basicBlockLabel,
             PVector<LlvmInstruction> instructions,
             PMap<PatioLocalRef, LlvmType> patioLocalToType,
@@ -1308,8 +1348,7 @@ public class LlvmPatioCodeGen {
             NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator,
             PStack<LoopLabels> loopLabels
         ) {
-            this.functionNameToReturnType = functionNameToReturnType;
-            this.exceptionTypeNameToTag = exceptionTypeNameToTag;
+            this.moduleContext = moduleContext;
             this.basicBlocks = basicBlocks;
             this.basicBlockLabel = basicBlockLabel;
             this.instructions = instructions;
@@ -1320,17 +1359,16 @@ public class LlvmPatioCodeGen {
         }
 
         private LlvmType functionNameToReturnType(String functionName) {
-            return this.functionNameToReturnType.get(functionName);
+            return this.moduleContext.functionNameToReturnType(functionName);
         }
 
         public int exceptionTag(String exceptionTypeName) {
-            return exceptionTypeNameToTag.get(exceptionTypeName);
+            return this.moduleContext.exceptionTag(exceptionTypeName);
         }
 
         private FunctionContext addTerminatorInstruction(LlvmTerminatorInstruction terminatorInstruction) {
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks.plus(new LlvmBasicBlock(
                     this.basicBlockLabel,
                     this.instructions,
@@ -1358,8 +1396,7 @@ public class LlvmPatioCodeGen {
                 );
 
             return new FunctionContext(
-                context.functionNameToReturnType,
-                context.exceptionTypeNameToTag,
+                context.moduleContext,
                 context.basicBlocks,
                 Optional.of(basicBlockLabel.name()),
                 context.instructions,
@@ -1389,8 +1426,7 @@ public class LlvmPatioCodeGen {
 
         public FunctionContext addPatioLocalToType(PatioLocalRef patioLocal, LlvmType llvmType) {
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1403,8 +1439,7 @@ public class LlvmPatioCodeGen {
 
         public FunctionContext addInstruction(LlvmInstruction instruction) {
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions.plus(instruction),
@@ -1419,8 +1454,7 @@ public class LlvmPatioCodeGen {
             var result = localNameAllocator.allocateName(space, Optional.empty());
 
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1435,8 +1469,7 @@ public class LlvmPatioCodeGen {
             var result = localNameAllocator.allocateName(space, Optional.of(name));
 
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1455,8 +1488,7 @@ public class LlvmPatioCodeGen {
 
         public FunctionContext enterLoop(LlvmVariableLocal startLabel, LlvmVariableLocal endLabel) {
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1469,8 +1501,7 @@ public class LlvmPatioCodeGen {
 
         public FunctionContext exitLoop() {
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,
@@ -1491,8 +1522,7 @@ public class LlvmPatioCodeGen {
 
         public FunctionContext addPatioLocal(PatioLocalRef patioLocal, LlvmOperand llvmOperand) {
             return new FunctionContext(
-                this.functionNameToReturnType,
-                this.exceptionTypeNameToTag,
+                this.moduleContext,
                 this.basicBlocks,
                 this.basicBlockLabel,
                 this.instructions,