Push unit at end of function definition body

3 files changed, 27 insertions, 10 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()
         );
 
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
index 1fa6c41a..1c90a0a7 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
@@ -36,7 +36,8 @@ public class StackCompilerFunctionDefinitionTests {
                 has("freeVariables", x -> x.freeVariables(), containsExactly(equalTo("z"))),
                 has("body", x -> x.body(), isSequence(
                     isPushStringInstruction("hello"),
-                    isDiscardInstruction()
+                    isDiscardInstruction(),
+                    isPushUnitInstruction()
                 ))
             ),
             isLocalStoreInstruction("f")
diff --git a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
index bfac50ce..0cbad43d 100644
--- a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
+++ b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
@@ -73,4 +73,10 @@ public class InstructionMatchers {
             has("value", x -> x.value(), equalTo(value))
         );
     }
+
+    public static Matcher<Instruction> isPushUnitInstruction() {
+        return instanceOf(
+            PushUnitInstruction.class
+        );
+    }
 }