From aa7ae2695c1ee5b919832b84df401614a7500a27 Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Sun, 11 Jun 2023 10:30:31 +0100
Subject: Push unit at end of function definition body

---
 .../shedlang/compiler/stackir/StackCompiler.java   | 28 +++++++++++++++-------
 1 file changed, 19 insertions(+), 9 deletions(-)

(limited to 'src/main/java')

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()
         );
 
-- 
cgit v1.2.3