Handle block returns properly

1 files changed, 26 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index d3b8801c..a8bc9d45 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -4,6 +4,7 @@ import org.pcollections.PVector;
 import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
+import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
 import java.util.stream.Collectors;
@@ -29,6 +30,10 @@ public class StackCompiler {
             });
     }
 
+    private static PVector<Instruction> compileBlockReturn(TypedBlockReturnNode blockReturn) {
+        return compileExpression(blockReturn.expression());
+    }
+
     private static PVector<Instruction> compileBoolLiteral(TypedBoolLiteralNode node) {
         return P.vector(new PushBoolInstruction(node.value(), node.source()));
     }
@@ -96,8 +101,24 @@ public class StackCompiler {
             .flatMap(statement -> compileFunctionStatement(statement).stream())
             .collect(Collectors.toList());
 
-        body.add(new PushUnitInstruction(node.source()));
-        body.add(new ReturnInstruction(node.source()));
+        var lastStatement = Lists.last(node.body());
+        // TODO: remove duplication with type checker (default to none behaviour).
+        // Extract block node?
+        var lastStatementBehaviour = lastStatement
+            .map(statement -> statement.behaviour())
+            .orElse(TypedFunctionStatementBehaviour.NONE);
+
+        switch (lastStatementBehaviour) {
+            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                body.add(new ReturnInstruction(node.source()));
+            }
+            case TypedFunctionStatementBehaviour.None none -> {
+                body.add(new PushUnitInstruction(node.source()));
+                body.add(new ReturnInstruction(node.source()));
+            }
+            case TypedFunctionStatementBehaviour.Terminate terminate -> {
+            }
+        }
 
         var defineFunctionInstruction = new DefineFunctionInstruction(
             node.name(),
@@ -120,6 +141,9 @@ public class StackCompiler {
     
     public static PVector<Instruction> compileFunctionStatement(TypedFunctionStatementNode node) {
         return switch (node) {
+            case TypedBlockReturnNode blockReturn ->
+                compileBlockReturn(blockReturn);
+            
             case TypedExpressionStatementNode expressionStatement ->
                 compileExpressionStatement(expressionStatement);