Add explicit return instruction to function bodies

3 files changed, 45 insertions, 11 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 2ad82eef..26319b43 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -1,6 +1,5 @@
 package org.shedlang.compiler.backends.llvm;
 
-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.LlvmParam;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmReturn;
@@ -64,8 +63,6 @@ public class LlvmCompiler {
                     functionContext = compileInstruction(bodyInstruction, functionContext);
                 }
                 // TODO: rejoin function context with context
-                // TODO: assert stack has exactly one element
-                var returnValue = functionContext.stack().get(0);
 
                 yield context.addTopLevelEntity(new LlvmFunctionDefinition(
                     topLevelEntityName,
@@ -77,13 +74,7 @@ public class LlvmCompiler {
                             "shed__env__0"
                         )
                     ),
-                    List.of(
-                        new LlvmBasicBlock(
-                            Optional.empty(),
-                            List.of(),
-                            new LlvmReturn(LlvmCompilerTypes.VALUE_TYPE, returnValue)
-                        )
-                    )
+                    functionContext.basicBlocks()
                 ));
             }
             case DiscardInstruction discardInstruction ->
@@ -129,6 +120,16 @@ public class LlvmCompiler {
 
             case PushUnitInstruction pushUnit ->
                 context.push(new LlvmOperandInt(0));
+
+            case ReturnInstruction returnInstruction -> {
+                var result = context.pop();
+                var value = result.value();
+                context = result.context();
+
+                yield context.addTerminatorInstruction(
+                    new LlvmReturn(LlvmCompilerTypes.VALUE_TYPE, value)
+                );
+            }
         };
     }
 }
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 d558f483..256079bc 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
@@ -1,9 +1,14 @@
 package org.shedlang.compiler.backends.llvm;
 
-import org.pcollections.*;
+import org.pcollections.Empty;
+import org.pcollections.PMap;
+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.LlvmTopLevelEntity;
 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.typechecker.TypeInfo;
@@ -16,6 +21,7 @@ public record LlvmCompilerContext(
     TypeInfo typeInfo,
     PStack<LlvmOperand> stack,
     PMap<String, LlvmOperand> locals,
+    PVector<LlvmBasicBlock> basicBlocks,
     PVector<LlvmInstruction> instructions,
     POrderedMap<String, LlvmTopLevelEntity> topLevelEntitiesByName,
     NameAllocator<LlvmCompilerLocalVariableSpace> localVariableNames
@@ -28,6 +34,7 @@ public record LlvmCompilerContext(
             Empty.stack(),
             Empty.map(),
             Empty.vector(),
+            Empty.vector(),
             POrderedMap.empty(),
             NameAllocator.empty()
         );
@@ -38,6 +45,7 @@ public record LlvmCompilerContext(
             typeInfo,
             stack.plus(operand),
             locals,
+            basicBlocks,
             instructions,
             topLevelEntitiesByName,
             localVariableNames
@@ -49,6 +57,7 @@ public record LlvmCompilerContext(
             typeInfo,
             stack.minus(0),
             locals,
+            basicBlocks,
             instructions,
             topLevelEntitiesByName,
             localVariableNames
@@ -60,6 +69,7 @@ public record LlvmCompilerContext(
             typeInfo,
             stack,
             locals.plus(name, value),
+            basicBlocks,
             instructions,
             topLevelEntitiesByName,
             localVariableNames
@@ -75,12 +85,31 @@ public record LlvmCompilerContext(
             typeInfo,
             stack,
             locals,
+            basicBlocks,
             instructions.plus(instruction),
             topLevelEntitiesByName,
             localVariableNames
         );
     }
 
+    public LlvmCompilerContext addTerminatorInstruction(
+        LlvmTerminatorInstruction terminatorInstruction
+    ) {
+        return new LlvmCompilerContext(
+            typeInfo,
+            stack,
+            locals,
+            basicBlocks.plus(new LlvmBasicBlock(
+                Optional.empty(),
+                instructions,
+                terminatorInstruction
+            )),
+            Empty.vector(),
+            topLevelEntitiesByName,
+            localVariableNames
+        );
+    }
+
     public LlvmCompilerResult<LlvmGlobalDefinition> addGlobalDefinition(
         LlvmGlobalDefinitionBuilder builder
     ) {
@@ -101,6 +130,7 @@ public record LlvmCompilerContext(
             typeInfo,
             stack,
             locals,
+            basicBlocks,
             instructions,
             topLevelEntitiesByName.plus(topLevelEntity.name(), topLevelEntity),
             localVariableNames
@@ -152,6 +182,7 @@ public record LlvmCompilerContext(
             typeInfo,
             stack,
             locals,
+            basicBlocks,
             instructions,
             topLevelEntitiesByName,
             result.allocator()
@@ -168,6 +199,7 @@ public record LlvmCompilerContext(
             Empty.stack(),
             Empty.map(),
             Empty.vector(),
+            Empty.vector(),
             topLevelEntitiesByName,
             NameAllocator.empty()
         );
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 6b8452c1..b8356ab4 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -76,6 +76,7 @@ public class StackCompiler {
             .collect(Collectors.toList());
 
         body.add(new PushUnitInstruction(node.source()));
+        body.add(new ReturnInstruction(node.source()));
 
         var defineFunctionInstruction = new DefineFunctionInstruction(
             node.name(),