Push unit at end of function definition body

1 files changed, 19 insertions, 9 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index aa59bc81..ff9b7de2 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -5,6 +5,8 @@ import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.util.P;
 
+import java.util.stream.Collectors;
+
 public class StackCompiler {
     private StackCompiler() {
     }
@@ -61,18 +63,26 @@ public class StackCompiler {
     }
 
     private static PVector<Instruction> compileFunctionDefinition(TypedFunctionDefinitionNode node) {
+        var positionalParams = node.params().positional().stream()
+            .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
+            .toList();
+
+        var namedParams = node.params().named().stream()
+            .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
+            .toList();
+
+        var body = node.body().stream()
+            .flatMap(statement -> compileFunctionStatement(statement).stream())
+            .collect(Collectors.toList());
+
+        body.add(new PushUnitInstruction(node.source()));
+
         var defineFunctionInstruction = new DefineFunctionInstruction(
             node.name(),
-            node.params().positional().stream()
-                .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
-                .toList(),
-            node.params().named().stream()
-                .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
-                .toList(),
+            positionalParams,
+            namedParams,
             node.freeVariables(),
-            node.body().stream()
-                .flatMap(statement -> compileFunctionStatement(statement).stream())
-                .toList(),
+            body,
             node.source()
         );