Preserve top level entities when exiting function

2 files changed, 19 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
index 26319b43..4ee42fa8 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -62,7 +62,9 @@ public class LlvmCompiler {
                 for (var bodyInstruction : defineFunction.body()) {
                     functionContext = compileInstruction(bodyInstruction, functionContext);
                 }
-                // TODO: rejoin function context with context
+                var bodyResult = context.exitFunction(functionContext);
+                var body = bodyResult.value();
+                context = bodyResult.context();
 
                 yield context.addTopLevelEntity(new LlvmFunctionDefinition(
                     topLevelEntityName,
@@ -74,7 +76,7 @@ public class LlvmCompiler {
                             "shed__env__0"
                         )
                     ),
-                    functionContext.basicBlocks()
+                    body
                 ));
             }
             case DiscardInstruction discardInstruction ->
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
index 256079bc..d7d8c84d 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
@@ -15,6 +15,7 @@ import org.shedlang.compiler.typechecker.TypeInfo;
 import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Optional;
 
 public record LlvmCompilerContext(
@@ -204,4 +205,18 @@ public record LlvmCompilerContext(
             NameAllocator.empty()
         );
     }
+
+    public LlvmCompilerResult<List<LlvmBasicBlock>> exitFunction(
+        LlvmCompilerContext functionContext
+    ) {
+        return new LlvmCompilerContext(
+            typeInfo,
+            stack,
+            locals,
+            basicBlocks,
+            instructions,
+            topLevelEntitiesByName.plusAll(functionContext.topLevelEntitiesByName),
+            localVariableNames
+        ).toResult(functionContext.basicBlocks);
+    }
 }