Split out LLVM code gen contexts

6 files changed, 454 insertions, 416 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 73b0f0a7..2c40bc78 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.backends.llvm;
 
 import org.shedlang.compiler.backends.Backend;
 import org.shedlang.compiler.backends.llvm.codegen.LlvmPatioCodeGen;
+import org.shedlang.compiler.backends.llvm.codegen.LlvmPatioCodeGenGlobalContext;
 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.LlvmModule;
@@ -50,7 +51,7 @@ public class LlvmBackend implements Backend {
         TargetDir targetDir
     ) throws IOException, InterruptedException {
         var objectFilePaths = new ArrayList<Path>();
-        var globalContext = LlvmPatioCodeGen.GlobalContext.EMPTY;
+        var globalContext = LlvmPatioCodeGenGlobalContext.EMPTY;
 
         for (var namespace : namespaces) {
             var patioCompilationUnit = PatioCodeGen.generateNamespace(
@@ -103,7 +104,7 @@ public class LlvmBackend implements Backend {
     private LlvmModule generateMain(
         NamespaceName main,
         TypeInfos typeInfos,
-        LlvmPatioCodeGen.GlobalContext globalContext
+        LlvmPatioCodeGenGlobalContext globalContext
     ) {
         // TODO: use more appropriate source
         var source = Sources.NULL;
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java
index 8acf6c34..20c2680a 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java
@@ -1,10 +1,7 @@
 package org.shedlang.compiler.backends.llvm.codegen;
 
-import org.pcollections.PMap;
-import org.pcollections.PStack;
 import org.pcollections.PVector;
 import org.pcollections.TreePVector;
-import org.shedlang.compiler.backends.NameAllocator;
 import org.shedlang.compiler.backends.llvm.LlvmCodeGenLocalVariableSpace;
 import org.shedlang.compiler.backends.llvm.ir.ast.*;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.*;
@@ -79,11 +76,11 @@ public class LlvmPatioCodeGen {
         )
     );
 
-    public static GlobalContext.Result<LlvmModule> generateCompilationUnit(
+    public static LlvmPatioCodeGenGlobalContext.Result<LlvmModule> generateCompilationUnit(
         PatioCompilationUnit patioCompilationUnit,
-        GlobalContext globalContext
+        LlvmPatioCodeGenGlobalContext globalContext
     ) {
-        var context = ModuleContext.EMPTY
+        var context = LlvmPatioCodeGenModuleContext.EMPTY
             .addTopLevelEntity(ALLOC_DECLARATION)
             .addTopLevelEntity(SETJMP_DECLARATION)
             .addTopLevelEntity(EXCEPTION_HANDLER_DECLARATION)
@@ -122,14 +119,12 @@ public class LlvmPatioCodeGen {
             context = generateCompilationUnitElement(patioCompilationUnitElement, context);
         }
 
-        var module = context.toModule();
-
-        return context.globalContext.toResult(module);
+        return context.toModule();
     }
 
-    private static ModuleContext generateCompilationUnitElement(
+    private static LlvmPatioCodeGenModuleContext generateCompilationUnitElement(
         PatioCompilationUnitElement patioCompilationUnitElement,
-        ModuleContext context
+        LlvmPatioCodeGenModuleContext context
     ) {
         return switch (patioCompilationUnitElement) {
             case PatioDataElementImmutable patioDataElementImmutable ->
@@ -149,9 +144,9 @@ public class LlvmPatioCodeGen {
         };
     }
 
-    private static ModuleContext generateDataElementImmutable(
+    private static LlvmPatioCodeGenModuleContext generateDataElementImmutable(
         PatioDataElementImmutable patioDataElementImmutable,
-        ModuleContext context
+        LlvmPatioCodeGenModuleContext context
     ) {
         return context.addTopLevelEntity(new LlvmGlobalDefinition(
             patioDataElementImmutable.name(),
@@ -163,9 +158,9 @@ public class LlvmPatioCodeGen {
         ));
     }
 
-    private static ModuleContext generateDataElementMutable(
+    private static LlvmPatioCodeGenModuleContext generateDataElementMutable(
         PatioDataElementMutable patioDataElementMutable,
-        ModuleContext context
+        LlvmPatioCodeGenModuleContext context
     ) {
         return context.addTopLevelEntity(new LlvmGlobalDefinition(
             patioDataElementMutable.name(),
@@ -177,9 +172,9 @@ public class LlvmPatioCodeGen {
         ));
     }
 
-    private static ModuleContext generateFunctionDeclaration(
+    private static LlvmPatioCodeGenModuleContext generateFunctionDeclaration(
         PatioFunctionDeclaration patioFunctionDeclaration,
-        ModuleContext context
+        LlvmPatioCodeGenModuleContext context
     ) {
         return context.addTopLevelEntity(new LlvmFunctionDeclaration(
             patioFunctionDeclaration.name(),
@@ -193,11 +188,11 @@ public class LlvmPatioCodeGen {
         ));
     }
 
-    private static ModuleContext generateFunctionDefinition(
+    private static LlvmPatioCodeGenModuleContext generateFunctionDefinition(
         PatioFunctionDefinition patioFunctionDefinition,
-        ModuleContext context
+        LlvmPatioCodeGenModuleContext moduleContext
     ) {
-        var functionContext = context.enterFunction();
+        var functionContext = moduleContext.enterFunction();
 
         for (var patioLocal : patioFunctionDefinition.locals()) {
             functionContext = functionContext.addPatioLocalToType(
@@ -214,9 +209,12 @@ public class LlvmPatioCodeGen {
         }
 
         functionContext = generateInstructions(patioFunctionDefinition.body(), functionContext);
-        var llvmBody = functionContext.basicBlocks;
 
-        return context.addTopLevelEntity(new LlvmFunctionDefinition(
+        var llvmBodyResult = functionContext.toBasicBlocks();
+        var llvmBody = llvmBodyResult.value();
+        moduleContext = llvmBodyResult.context();
+
+        return moduleContext.addTopLevelEntity(new LlvmFunctionDefinition(
             patioFunctionDefinition.name(),
             patioTypeToLlvmType(patioFunctionDefinition.returnType()),
             StreamSupport.stream(patioFunctionDefinition.params().spliterator(), false)
@@ -229,9 +227,9 @@ public class LlvmPatioCodeGen {
         ));
     }
 
-    private static FunctionContext generateInstructions(
+    private static LlvmPatioCodeGenFunctionContext generateInstructions(
         List<PatioInstruction> patioInstructions,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         for (var patioInstruction : patioInstructions) {
             context = generateInstruction(patioInstruction, context);
@@ -240,19 +238,19 @@ public class LlvmPatioCodeGen {
         return context;
     }
 
-    private static FunctionContext generateInstruction(
+    private static LlvmPatioCodeGenFunctionContext generateInstruction(
         PatioInstruction patioInstruction,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         return switch (patioInstruction) {
             case PatioAllocHeap patioAllocHeap -> {
                 var sizeResult = sizeOf(patioDataLayoutToLlvmType(patioAllocHeap.layout()), context);
-                var size = sizeResult.value;
-                context = sizeResult.context;
+                var size = sizeResult.value();
+                context = sizeResult.context();
 
                 var targetResult = context.addPatioLocalToLlvmLocal(patioAllocHeap.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(ALLOC_DECLARATION.call(
                     Optional.of(target),
@@ -262,8 +260,8 @@ public class LlvmPatioCodeGen {
 
             case PatioAllocStack patioAllocStack -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioAllocStack.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmAlloca(
                     target,
@@ -273,8 +271,8 @@ public class LlvmPatioCodeGen {
 
             case PatioBoolToInt patioBoolToInt -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioBoolToInt.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmZext(
                     target,
@@ -288,8 +286,8 @@ public class LlvmPatioCodeGen {
                 var target = Optional.<LlvmVariable>empty();
                 if (patioCall.target().isPresent()) {
                     var targetResult = context.addPatioLocalToLlvmLocal(patioCall.target().get());
-                    target = Optional.of(targetResult.value);
-                    context = targetResult.context;
+                    target = Optional.of(targetResult.value());
+                    context = targetResult.context();
                 }
 
                 var c = context;
@@ -306,8 +304,8 @@ public class LlvmPatioCodeGen {
 
             case PatioCallPointer patioCallPointer -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioCallPointer.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 var c = context;
 
@@ -323,8 +321,8 @@ public class LlvmPatioCodeGen {
 
             case PatioDataPtrToInt patioDataPtrToInt -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioDataPtrToInt.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmPtrToInt(
                     target,
@@ -385,8 +383,8 @@ public class LlvmPatioCodeGen {
 
             case PatioIntAdd patioIntAdd -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioIntAdd.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmAdd(
                     target,
@@ -398,8 +396,8 @@ public class LlvmPatioCodeGen {
 
             case PatioIntEquals patioIntEquals -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioIntEquals.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmIcmp(
                     target,
@@ -412,8 +410,8 @@ public class LlvmPatioCodeGen {
 
             case PatioIntNotEqual patioIntNotEqual -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioIntNotEqual.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmIcmp(
                     target,
@@ -426,8 +424,8 @@ public class LlvmPatioCodeGen {
 
             case PatioIntSub patioIntSub -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioIntSub.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmSub(
                     target,
@@ -439,8 +437,8 @@ public class LlvmPatioCodeGen {
 
             case PatioIntToBool patioIntToBool -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioIntToBool.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmTrunc(
                     target,
@@ -452,8 +450,8 @@ public class LlvmPatioCodeGen {
 
             case PatioIntToDataPtr patioIntToDataPtr -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioIntToDataPtr.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 yield context.addInstruction(new LlvmIntToPtr(
                     target,
@@ -500,12 +498,12 @@ public class LlvmPatioCodeGen {
 
             case PatioMemoryLoad patioMemoryLoad -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioMemoryLoad.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 var pointerResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TEMPORARY);
-                var pointer = pointerResult.value;
-                context = pointerResult.context;
+                var pointer = pointerResult.value();
+                context = pointerResult.context();
 
                 context = context.addInstruction(new LlvmGetElementPtr(
                     pointer,
@@ -525,8 +523,8 @@ public class LlvmPatioCodeGen {
 
             case PatioMemoryStore patioMemoryStore -> {
                 var pointerResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TEMPORARY);
-                var pointer = pointerResult.value;
-                context = pointerResult.context;
+                var pointer = pointerResult.value();
+                context = pointerResult.context();
 
                 context = context.addInstruction(new LlvmGetElementPtr(
                     pointer,
@@ -647,8 +645,8 @@ public class LlvmPatioCodeGen {
 
             case PatioTry try_ -> {
                 var jmpBufResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.JMP_BUF);
-                var jmpBuf = jmpBufResult.value;
-                context = jmpBufResult.context;
+                var jmpBuf = jmpBufResult.value();
+                context = jmpBufResult.context();
 
                 context = context.addInstruction(new LlvmAlloca(
                     jmpBuf,
@@ -656,8 +654,8 @@ public class LlvmPatioCodeGen {
                 ));
 
                 var previousExceptionHandlerResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TRY_PREVIOUS_EXCEPTION_HANDLER);
-                var previousExceptionHandler = previousExceptionHandlerResult.value;
-                context = previousExceptionHandlerResult.context;
+                var previousExceptionHandler = previousExceptionHandlerResult.value();
+                context = previousExceptionHandlerResult.context();
 
                 context = context.addInstruction(new LlvmLoad(
                     previousExceptionHandler,
@@ -678,8 +676,8 @@ public class LlvmPatioCodeGen {
                 ));
 
                 var setjmpValueResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.SETJMP_VALUE);
-                var setjmpValue = setjmpValueResult.value;
-                context = setjmpValueResult.context;
+                var setjmpValue = setjmpValueResult.value();
+                context = setjmpValueResult.context();
 
                 context = context.addInstruction(SETJMP_DECLARATION.call(
                     Optional.of(setjmpValue),
@@ -730,8 +728,8 @@ public class LlvmPatioCodeGen {
                     context = catchBodyEndLabelResult.context();
 
                     var exceptionValueTargetResult = context.addPatioLocalToLlvmLocal(catchClause.target());
-                    var exceptionValueTarget = exceptionValueTargetResult.value;
-                    context = exceptionValueTargetResult.context;
+                    var exceptionValueTarget = exceptionValueTargetResult.value();
+                    context = exceptionValueTargetResult.context();
 
                     branches.add(new Branch(
                         catchBodyStartLabel,
@@ -752,8 +750,8 @@ public class LlvmPatioCodeGen {
                 }
 
                 var rethrowStartLabelResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.TRY_RETHROW_START);
-                var rethrowStartLabel = rethrowStartLabelResult.value;
-                context = rethrowStartLabelResult.context;
+                var rethrowStartLabel = rethrowStartLabelResult.value();
+                context = rethrowStartLabelResult.context();
 
                 context = context.addTerminatorInstruction(new LlvmSwitch(
                     LlvmTypes.INT_32,
@@ -784,8 +782,8 @@ public class LlvmPatioCodeGen {
 
             case PatioVarLoad patioVarLoad -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioVarLoad.target());
-                var target = targetResult.value;
-                context = targetResult.context;
+                var target = targetResult.value();
+                context = targetResult.context();
 
                 context = context.addInstruction(new LlvmLoad(
                     target,
@@ -799,8 +797,8 @@ public class LlvmPatioCodeGen {
             case PatioVarStore patioVarStore -> {
                 if (!context.patioLocalHasOperand(patioVarStore.var())) {
                     var varResult = context.addPatioLocalToLlvmLocal(patioVarStore.var());
-                    var var = varResult.value;
-                    context = varResult.context;
+                    var var = varResult.value();
+                    context = varResult.context();
 
                     context = context.addInstruction(new LlvmAlloca(
                         var,
@@ -835,11 +833,11 @@ public class LlvmPatioCodeGen {
         }
     }
 
-    private static FunctionContext generateBranches(
+    private static LlvmPatioCodeGenFunctionContext generateBranches(
         List<PatioLocalRef> targets,
         List<Branch> branches,
         LlvmVariableLocal endLabel,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         var phiPredecessors = new ArrayList<PhiPredecessor>();
         for (var branch : branches) {
@@ -880,17 +878,17 @@ public class LlvmPatioCodeGen {
 
     record PhiPredecessor(PatioBlock block, LlvmVariableLocal label) {}
 
-    private static FunctionContext generatePhiInstructions(
+    private static LlvmPatioCodeGenFunctionContext generatePhiInstructions(
         List<PatioLocalRef> targets,
         List<PhiPredecessor> predecessors,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         for (var targetIndex = 0; targetIndex < targets.size(); targetIndex++) {
             var patioTarget = targets.get(targetIndex);
 
             var llvmTargetResult = context.addPatioLocalToLlvmLocal(patioTarget);
-            var llvmTarget = llvmTargetResult.value;
-            context = llvmTargetResult.context;
+            var llvmTarget = llvmTargetResult.value();
+            context = llvmTargetResult.context();
 
             var blockTargetIndex = targetIndex;
             var blockContext = context;
@@ -1007,7 +1005,7 @@ public class LlvmPatioCodeGen {
 
     private static LlvmTypedOperand patioOperandToLlvmTypedOperand(
         PatioOperand patioOperand,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         return new LlvmTypedOperand(
             patioOperandToLlvmType(patioOperand, context),
@@ -1017,7 +1015,7 @@ public class LlvmPatioCodeGen {
 
     private static LlvmType patioOperandToLlvmType(
         PatioOperand patioOperand,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         return switch (patioOperand) {
             case PatioBool patioBool ->
@@ -1039,7 +1037,7 @@ public class LlvmPatioCodeGen {
 
     private static LlvmOperand patioOperandToLlvmOperand(
         PatioOperand patioOperand,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         return switch (patioOperand) {
             case PatioBool patioBool ->
@@ -1182,9 +1180,9 @@ public class LlvmPatioCodeGen {
         };
     }
 
-    public static FunctionContext.Result<LlvmOperand> sizeOf(
+    public static LlvmPatioCodeGenFunctionContext.Result<LlvmOperand> sizeOf(
         LlvmType type,
-        FunctionContext context
+        LlvmPatioCodeGenFunctionContext context
     ) {
         var sizePtrResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.SIZE_PTR);
         var sizePtr = sizePtrResult.value();
@@ -1212,335 +1210,4 @@ 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 = 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 ModuleContext(
-            GlobalContext globalContext,
-            PMap<String, LlvmType> functionNameToReturnType,
-            PVector<LlvmTopLevelEntity> topLevelEntities
-        ) {
-            this.globalContext = globalContext;
-            this.functionNameToReturnType = functionNameToReturnType;
-            this.topLevelEntities = topLevelEntities;
-        }
-
-        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
-            );
-        }
-
-        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)
-            );
-        }
-
-        public ModuleContext allocateExceptionTag(String exceptionTypeName) {
-            return new ModuleContext(
-                this.globalContext.allocateExceptionTag(exceptionTypeName),
-                this.functionNameToReturnType,
-                this.topLevelEntities
-            );
-        }
-
-        public int exceptionTag(String exceptionTypeName) {
-            return this.globalContext.exceptionTag(exceptionTypeName);
-        }
-
-        public FunctionContext enterFunction() {
-            return FunctionContext.of(this);
-        }
-
-        private <T> Result<T> toResult(T value) {
-            return new Result<>(value, this);
-        }
-
-        public record Result<T>(T value, ModuleContext context) {
-        }
-    }
-
-    private static class FunctionContext {
-        private static FunctionContext of(ModuleContext moduleContext) {
-            return new FunctionContext(
-                moduleContext,
-                P.vector(),
-                Optional.empty(),
-                P.vector(),
-                P.map(),
-                P.map(),
-                NameAllocator.empty(),
-                P.stack()
-            );
-        }
-
-        private final ModuleContext moduleContext;
-        private final PVector<LlvmBasicBlock> basicBlocks;
-        private final Optional<String> basicBlockLabel;
-        private final PVector<LlvmInstruction> instructions;
-        private final PMap<PatioLocalRef, LlvmType> patioLocalToType;
-        private final PMap<PatioLocalRef, LlvmOperand> patioLocalToOperand;
-        private final NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator;
-        private final PStack<LoopLabels> loopLabels;
-
-        private FunctionContext(
-            ModuleContext moduleContext,
-            PVector<LlvmBasicBlock> basicBlocks,
-            Optional<String> basicBlockLabel,
-            PVector<LlvmInstruction> instructions,
-            PMap<PatioLocalRef, LlvmType> patioLocalToType,
-            PMap<PatioLocalRef, LlvmOperand> patioLocalToOperand,
-            NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator,
-            PStack<LoopLabels> loopLabels
-        ) {
-            this.moduleContext = moduleContext;
-            this.basicBlocks = basicBlocks;
-            this.basicBlockLabel = basicBlockLabel;
-            this.instructions = instructions;
-            this.patioLocalToType = patioLocalToType;
-            this.patioLocalToOperand = patioLocalToOperand;
-            this.localNameAllocator = localNameAllocator;
-            this.loopLabels = loopLabels;
-        }
-
-        private LlvmType functionNameToReturnType(String functionName) {
-            return this.moduleContext.functionNameToReturnType(functionName);
-        }
-
-        public int exceptionTag(String exceptionTypeName) {
-            return this.moduleContext.exceptionTag(exceptionTypeName);
-        }
-
-        private FunctionContext addTerminatorInstruction(LlvmTerminatorInstruction terminatorInstruction) {
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks.plus(new LlvmBasicBlock(
-                    this.basicBlockLabel,
-                    this.instructions,
-                    terminatorInstruction
-                )),
-                Optional.empty(),
-                P.vector(),
-                this.patioLocalToType,
-                this.patioLocalToOperand,
-                this.localNameAllocator,
-                this.loopLabels
-            );
-        }
-
-        /**
-         * Start a new basic block with the given label. If already in a basic
-         * block, then the existing basic block is terminated with an
-         * unconditional jump to the new basic block.
-         */
-        public FunctionContext startBasicBlock(LlvmVariableLocal basicBlockLabel) {
-            var context = this.basicBlockLabel.isEmpty() && this.instructions.isEmpty()
-                ? this
-                : this.addTerminatorInstruction(
-                    new LlvmBrUnconditional(basicBlockLabel)
-                );
-
-            return new FunctionContext(
-                context.moduleContext,
-                context.basicBlocks,
-                Optional.of(basicBlockLabel.name()),
-                context.instructions,
-                context.patioLocalToType,
-                context.patioLocalToOperand,
-                context.localNameAllocator,
-                context.loopLabels
-            );
-        }
-
-        public LlvmType patioLocalToType(PatioLocalRef value) {
-            return this.patioLocalToType.get(value);
-        }
-
-        public LlvmOperand patioLocalToOperand(PatioLocalRef value) {
-            var operand = this.patioLocalToOperand.get(value);
-            if (operand == null) {
-                throw new RuntimeException("no operand for Patio local: " + value.name());
-            } else {
-                return operand;
-            }
-        }
-
-        public boolean patioLocalHasOperand(PatioLocalRef value) {
-            return this.patioLocalToOperand.containsKey(value);
-        }
-
-        public FunctionContext addPatioLocalToType(PatioLocalRef patioLocal, LlvmType llvmType) {
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions,
-                this.patioLocalToType.plus(patioLocal, llvmType),
-                this.patioLocalToOperand,
-                this.localNameAllocator,
-                this.loopLabels
-            );
-        }
-
-        public FunctionContext addInstruction(LlvmInstruction instruction) {
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions.plus(instruction),
-                this.patioLocalToType,
-                this.patioLocalToOperand,
-                this.localNameAllocator,
-                this.loopLabels
-            );
-        }
-
-        public Result<LlvmVariableLocal> allocateLocalName(LlvmCodeGenLocalVariableSpace space) {
-            var result = localNameAllocator.allocateName(space, Optional.empty());
-
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions,
-                this.patioLocalToType,
-                this.patioLocalToOperand,
-                result.allocator(),
-                this.loopLabels
-            ).toResult(new LlvmVariableLocal(result.name()));
-        }
-
-        public Result<LlvmVariableLocal> allocateLocalName(LlvmCodeGenLocalVariableSpace space, String name) {
-            var result = localNameAllocator.allocateName(space, Optional.of(name));
-
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions,
-                this.patioLocalToType,
-                this.patioLocalToOperand,
-                result.allocator(),
-                this.loopLabels
-            ).toResult(new LlvmVariableLocal(result.name()));
-        }
-
-        public Result<LlvmVariableLocal> addPatioLocalToLlvmLocal(PatioLocalRef patioLocal) {
-            var result = allocateLocalName(LlvmCodeGenLocalVariableSpace.LOCAL, patioLocal.name());
-
-            return result.context.addPatioLocal(patioLocal, result.value).toResult(result.value);
-        }
-
-        public FunctionContext enterLoop(LlvmVariableLocal startLabel, LlvmVariableLocal endLabel) {
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions,
-                this.patioLocalToType,
-                this.patioLocalToOperand,
-                this.localNameAllocator,
-                this.loopLabels.plus(new LoopLabels(startLabel, endLabel))
-            );
-        }
-
-        public FunctionContext exitLoop() {
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions,
-                this.patioLocalToType,
-                this.patioLocalToOperand,
-                this.localNameAllocator,
-                this.loopLabels.minus(0)
-            );
-        }
-
-        public LlvmVariableLocal loopStartLabel() {
-            return this.loopLabels.get(0).start;
-        }
-
-        public LlvmVariableLocal loopEndLabel() {
-            return this.loopLabels.get(0).end;
-        }
-
-        public FunctionContext addPatioLocal(PatioLocalRef patioLocal, LlvmOperand llvmOperand) {
-            return new FunctionContext(
-                this.moduleContext,
-                this.basicBlocks,
-                this.basicBlockLabel,
-                this.instructions,
-                this.patioLocalToType,
-                this.patioLocalToOperand.plus(patioLocal, llvmOperand),
-                this.localNameAllocator,
-                this.loopLabels
-            );
-        }
-
-        private <T> Result<T> toResult(T value) {
-            return new Result<>(value, this);
-        }
-
-        public record Result<T>(T value, FunctionContext context) {
-        }
-    }
-
-    private record LoopLabels(LlvmVariableLocal start, LlvmVariableLocal end) {}
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenFunctionContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenFunctionContext.java
new file mode 100644
index 00000000..1ab4b78a
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenFunctionContext.java
@@ -0,0 +1,250 @@
+package org.shedlang.compiler.backends.llvm.codegen;
+
+import org.pcollections.PMap;
+import org.pcollections.PStack;
+import org.pcollections.PVector;
+import org.shedlang.compiler.backends.NameAllocator;
+import org.shedlang.compiler.backends.llvm.LlvmCodeGenLocalVariableSpace;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmBrUnconditional;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmInstruction;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmTerminatorInstruction;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmType;
+import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
+import org.shedlang.compiler.util.P;
+
+import java.util.Optional;
+
+class LlvmPatioCodeGenFunctionContext {
+    static LlvmPatioCodeGenFunctionContext of(LlvmPatioCodeGenModuleContext moduleContext) {
+        return new LlvmPatioCodeGenFunctionContext(
+            moduleContext,
+            P.vector(),
+            Optional.empty(),
+            P.vector(),
+            P.map(),
+            P.map(),
+            NameAllocator.empty(),
+            P.stack()
+        );
+    }
+
+    private final LlvmPatioCodeGenModuleContext moduleContext;
+    private final PVector<LlvmBasicBlock> basicBlocks;
+    private final Optional<String> basicBlockLabel;
+    private final PVector<LlvmInstruction> instructions;
+    private final PMap<PatioLocalRef, LlvmType> patioLocalToType;
+    private final PMap<PatioLocalRef, LlvmOperand> patioLocalToOperand;
+    private final NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator;
+    private final PStack<LoopLabels> loopLabels;
+
+    private LlvmPatioCodeGenFunctionContext(
+        LlvmPatioCodeGenModuleContext moduleContext,
+        PVector<LlvmBasicBlock> basicBlocks,
+        Optional<String> basicBlockLabel,
+        PVector<LlvmInstruction> instructions,
+        PMap<PatioLocalRef, LlvmType> patioLocalToType,
+        PMap<PatioLocalRef, LlvmOperand> patioLocalToOperand,
+        NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator,
+        PStack<LoopLabels> loopLabels
+    ) {
+        this.moduleContext = moduleContext;
+        this.basicBlocks = basicBlocks;
+        this.basicBlockLabel = basicBlockLabel;
+        this.instructions = instructions;
+        this.patioLocalToType = patioLocalToType;
+        this.patioLocalToOperand = patioLocalToOperand;
+        this.localNameAllocator = localNameAllocator;
+        this.loopLabels = loopLabels;
+    }
+
+    LlvmType functionNameToReturnType(String functionName) {
+        return this.moduleContext.functionNameToReturnType(functionName);
+    }
+
+    int exceptionTag(String exceptionTypeName) {
+        return this.moduleContext.exceptionTag(exceptionTypeName);
+    }
+
+    LlvmPatioCodeGenModuleContext.Result<PVector<LlvmBasicBlock>> toBasicBlocks() {
+        // TODO: check there are no instructions without a terminator.
+
+        return this.moduleContext.toResult(this.basicBlocks);
+    }
+
+    LlvmPatioCodeGenFunctionContext addTerminatorInstruction(LlvmTerminatorInstruction terminatorInstruction) {
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks.plus(new LlvmBasicBlock(
+                this.basicBlockLabel,
+                this.instructions,
+                terminatorInstruction
+            )),
+            Optional.empty(),
+            P.vector(),
+            this.patioLocalToType,
+            this.patioLocalToOperand,
+            this.localNameAllocator,
+            this.loopLabels
+        );
+    }
+
+    /**
+     * Start a new basic block with the given label. If already in a basic
+     * block, then the existing basic block is terminated with an
+     * unconditional jump to the new basic block.
+     */
+    public LlvmPatioCodeGenFunctionContext startBasicBlock(LlvmVariableLocal basicBlockLabel) {
+        var context = this.basicBlockLabel.isEmpty() && this.instructions.isEmpty()
+            ? this
+            : this.addTerminatorInstruction(
+            new LlvmBrUnconditional(basicBlockLabel)
+        );
+
+        return new LlvmPatioCodeGenFunctionContext(
+            context.moduleContext,
+            context.basicBlocks,
+            Optional.of(basicBlockLabel.name()),
+            context.instructions,
+            context.patioLocalToType,
+            context.patioLocalToOperand,
+            context.localNameAllocator,
+            context.loopLabels
+        );
+    }
+
+    public LlvmType patioLocalToType(PatioLocalRef value) {
+        return this.patioLocalToType.get(value);
+    }
+
+    public LlvmOperand patioLocalToOperand(PatioLocalRef value) {
+        var operand = this.patioLocalToOperand.get(value);
+        if (operand == null) {
+            throw new RuntimeException("no operand for Patio local: " + value.name());
+        } else {
+            return operand;
+        }
+    }
+
+    public boolean patioLocalHasOperand(PatioLocalRef value) {
+        return this.patioLocalToOperand.containsKey(value);
+    }
+
+    public LlvmPatioCodeGenFunctionContext addPatioLocalToType(PatioLocalRef patioLocal, LlvmType llvmType) {
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions,
+            this.patioLocalToType.plus(patioLocal, llvmType),
+            this.patioLocalToOperand,
+            this.localNameAllocator,
+            this.loopLabels
+        );
+    }
+
+    public LlvmPatioCodeGenFunctionContext addInstruction(LlvmInstruction instruction) {
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions.plus(instruction),
+            this.patioLocalToType,
+            this.patioLocalToOperand,
+            this.localNameAllocator,
+            this.loopLabels
+        );
+    }
+
+    public Result<LlvmVariableLocal> allocateLocalName(LlvmCodeGenLocalVariableSpace space) {
+        var result = localNameAllocator.allocateName(space, Optional.empty());
+
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions,
+            this.patioLocalToType,
+            this.patioLocalToOperand,
+            result.allocator(),
+            this.loopLabels
+        ).toResult(new LlvmVariableLocal(result.name()));
+    }
+
+    public Result<LlvmVariableLocal> allocateLocalName(LlvmCodeGenLocalVariableSpace space, String name) {
+        var result = localNameAllocator.allocateName(space, Optional.of(name));
+
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions,
+            this.patioLocalToType,
+            this.patioLocalToOperand,
+            result.allocator(),
+            this.loopLabels
+        ).toResult(new LlvmVariableLocal(result.name()));
+    }
+
+    public Result<LlvmVariableLocal> addPatioLocalToLlvmLocal(PatioLocalRef patioLocal) {
+        var result = allocateLocalName(LlvmCodeGenLocalVariableSpace.LOCAL, patioLocal.name());
+
+        return result.context.addPatioLocal(patioLocal, result.value).toResult(result.value);
+    }
+
+    public LlvmPatioCodeGenFunctionContext enterLoop(LlvmVariableLocal startLabel, LlvmVariableLocal endLabel) {
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions,
+            this.patioLocalToType,
+            this.patioLocalToOperand,
+            this.localNameAllocator,
+            this.loopLabels.plus(new LoopLabels(startLabel, endLabel))
+        );
+    }
+
+    public LlvmPatioCodeGenFunctionContext exitLoop() {
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions,
+            this.patioLocalToType,
+            this.patioLocalToOperand,
+            this.localNameAllocator,
+            this.loopLabels.minus(0)
+        );
+    }
+
+    public LlvmVariableLocal loopStartLabel() {
+        return this.loopLabels.get(0).start();
+    }
+
+    public LlvmVariableLocal loopEndLabel() {
+        return this.loopLabels.get(0).end();
+    }
+
+    public LlvmPatioCodeGenFunctionContext addPatioLocal(PatioLocalRef patioLocal, LlvmOperand llvmOperand) {
+        return new LlvmPatioCodeGenFunctionContext(
+            this.moduleContext,
+            this.basicBlocks,
+            this.basicBlockLabel,
+            this.instructions,
+            this.patioLocalToType,
+            this.patioLocalToOperand.plus(patioLocal, llvmOperand),
+            this.localNameAllocator,
+            this.loopLabels
+        );
+    }
+
+    <T> Result<T> toResult(T value) {
+        return new Result<>(value, this);
+    }
+
+    public record Result<T>(T value, LlvmPatioCodeGenFunctionContext context) {
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenGlobalContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenGlobalContext.java
new file mode 100644
index 00000000..a0aacf7f
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenGlobalContext.java
@@ -0,0 +1,38 @@
+package org.shedlang.compiler.backends.llvm.codegen;
+
+import org.pcollections.PMap;
+import org.shedlang.compiler.util.P;
+
+public class LlvmPatioCodeGenGlobalContext {
+    public static final LlvmPatioCodeGenGlobalContext EMPTY = new LlvmPatioCodeGenGlobalContext(P.map());
+
+    private final PMap<String, Integer> exceptionTypeNameToTag;
+
+    private LlvmPatioCodeGenGlobalContext(
+        PMap<String, Integer> exceptionTypeNameToTag
+    ) {
+        this.exceptionTypeNameToTag = exceptionTypeNameToTag;
+    }
+
+    LlvmPatioCodeGenGlobalContext allocateExceptionTag(String exceptionTypeName) {
+        var tag = this.exceptionTypeNameToTag.size() + 1;
+        return new LlvmPatioCodeGenGlobalContext(
+            this.exceptionTypeNameToTag.plus(exceptionTypeName, tag)
+        );
+    }
+
+    int exceptionTag(String exceptionTypeName) {
+        return this.exceptionTypeNameToTag.get(exceptionTypeName);
+    }
+
+    LlvmPatioCodeGenModuleContext enterModule() {
+        return LlvmPatioCodeGenModuleContext.of(this);
+    }
+
+    <T> Result<T> toResult(T value) {
+        return new Result<>(value, this);
+    }
+
+    public record Result<T>(T value, LlvmPatioCodeGenGlobalContext context) {
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenModuleContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenModuleContext.java
new file mode 100644
index 00000000..ae630492
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGenModuleContext.java
@@ -0,0 +1,77 @@
+package org.shedlang.compiler.backends.llvm.codegen;
+
+import org.pcollections.PMap;
+import org.pcollections.PVector;
+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.types.LlvmType;
+import org.shedlang.compiler.util.P;
+
+class LlvmPatioCodeGenModuleContext {
+    static final LlvmPatioCodeGenModuleContext EMPTY = of(LlvmPatioCodeGenGlobalContext.EMPTY);
+
+    static LlvmPatioCodeGenModuleContext of(LlvmPatioCodeGenGlobalContext globalContext) {
+        return new LlvmPatioCodeGenModuleContext(globalContext, P.map(), P.vector());
+    }
+
+    private final LlvmPatioCodeGenGlobalContext globalContext;
+    private final PMap<String, LlvmType> functionNameToReturnType;
+    private final PVector<LlvmTopLevelEntity> topLevelEntities;
+
+    private LlvmPatioCodeGenModuleContext(
+        LlvmPatioCodeGenGlobalContext globalContext,
+        PMap<String, LlvmType> functionNameToReturnType,
+        PVector<LlvmTopLevelEntity> topLevelEntities
+    ) {
+        this.globalContext = globalContext;
+        this.functionNameToReturnType = functionNameToReturnType;
+        this.topLevelEntities = topLevelEntities;
+    }
+
+    LlvmPatioCodeGenGlobalContext.Result<LlvmModule> toModule() {
+        return this.globalContext.toResult(new LlvmModule(topLevelEntities));
+    }
+
+    LlvmPatioCodeGenModuleContext addFunctionNameToReturnType(String functionName, LlvmType returnType) {
+        return new LlvmPatioCodeGenModuleContext(
+            this.globalContext,
+            this.functionNameToReturnType.plus(functionName, returnType),
+            this.topLevelEntities
+        );
+    }
+
+    LlvmType functionNameToReturnType(String functionName) {
+        return this.functionNameToReturnType.get(functionName);
+    }
+
+    LlvmPatioCodeGenModuleContext addTopLevelEntity(LlvmTopLevelEntity topLevelEntity) {
+        return new LlvmPatioCodeGenModuleContext(
+            this.globalContext,
+            this.functionNameToReturnType,
+            this.topLevelEntities.plus(topLevelEntity)
+        );
+    }
+
+    LlvmPatioCodeGenModuleContext allocateExceptionTag(String exceptionTypeName) {
+        return new LlvmPatioCodeGenModuleContext(
+            this.globalContext.allocateExceptionTag(exceptionTypeName),
+            this.functionNameToReturnType,
+            this.topLevelEntities
+        );
+    }
+
+    int exceptionTag(String exceptionTypeName) {
+        return this.globalContext.exceptionTag(exceptionTypeName);
+    }
+
+    LlvmPatioCodeGenFunctionContext enterFunction() {
+        return LlvmPatioCodeGenFunctionContext.of(this);
+    }
+
+    <T> Result<T> toResult(T value) {
+        return new Result<>(value, this);
+    }
+
+    public record Result<T>(T value, LlvmPatioCodeGenModuleContext context) {
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LoopLabels.java b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LoopLabels.java
new file mode 100644
index 00000000..a2d7ef42
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LoopLabels.java
@@ -0,0 +1,5 @@
+package org.shedlang.compiler.backends.llvm.codegen;
+
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
+
+record LoopLabels(LlvmVariableLocal start, LlvmVariableLocal end) {}