Handle block returns properly

11 files changed, 124 insertions, 32 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockReturnNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockReturnNode.java
new file mode 100644
index 00000000..6bd5fb1d
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockReturnNode.java
@@ -0,0 +1,13 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+public record TypedBlockReturnNode(
+    TypedExpressionNode expression,
+    Source source
+) implements TypedFunctionStatementNode {
+    @Override
+    public TypedFunctionStatementBehaviour behaviour() {
+        return TypedFunctionStatementBehaviour.blockReturn(expression.type());
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionStatementNode.java
index 24aae178..25f1429d 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionStatementNode.java
@@ -6,4 +6,8 @@ public record TypedExpressionStatementNode(
     TypedExpressionNode expression,
     Source source
 ) implements TypedFunctionStatementNode {
+    @Override
+    public TypedFunctionStatementBehaviour behaviour() {
+        return TypedFunctionStatementBehaviour.NONE;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionReturnNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionReturnNode.java
index 883b5649..6f023338 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionReturnNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionReturnNode.java
@@ -6,4 +6,8 @@ public record TypedFunctionReturnNode(
     TypedExpressionNode expression,
     Source source
 ) implements TypedFunctionStatementNode {
+    @Override
+    public TypedFunctionStatementBehaviour behaviour() {
+        return TypedFunctionStatementBehaviour.TERMINATE;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementBehaviour.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementBehaviour.java
new file mode 100644
index 00000000..615634d2
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementBehaviour.java
@@ -0,0 +1,16 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.types.Type;
+
+public sealed interface TypedFunctionStatementBehaviour {
+    None NONE = new None();
+    Terminate TERMINATE = new Terminate();
+
+    static TypedFunctionStatementBehaviour blockReturn(Type type) {
+        return new BlockReturn(type);
+    }
+
+    record Terminate() implements TypedFunctionStatementBehaviour {}
+    record BlockReturn(Type type) implements TypedFunctionStatementBehaviour {}
+    record None() implements TypedFunctionStatementBehaviour {}
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
index bbc0fc6f..82fc3a8c 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
@@ -1,4 +1,5 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedExpressionStatementNode, TypedFunctionReturnNode, TypedValNode {
+public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedBlockReturnNode, TypedExpressionStatementNode, TypedFunctionReturnNode, TypedValNode {
+    TypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java
index f2e59ffe..04b9e1b7 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java
@@ -7,4 +7,8 @@ public record TypedValNode(
     TypedExpressionNode expression,
     Source source
 ) implements TypedFunctionStatementNode, TypedNamespaceStatementNode {
+    @Override
+    public TypedFunctionStatementBehaviour behaviour() {
+        return TypedFunctionStatementBehaviour.NONE;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java
new file mode 100644
index 00000000..55d6cca0
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java
@@ -0,0 +1,9 @@
+package org.shedlang.compiler.ast.untyped;
+
+import org.shedlang.compiler.sources.Source;
+
+public record UntypedBlockReturnNode(
+    UntypedExpressionNode expression,
+    Source source
+) implements UntypedFunctionStatementNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
index d970d6f3..c3e0466d 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.ast.untyped;
 
-public sealed interface UntypedFunctionStatementNode extends UntypedNode permits UntypedExpressionStatementNode, UntypedFunctionReturnNode, UntypedValNode {
+public sealed interface UntypedFunctionStatementNode extends UntypedNode permits UntypedBlockReturnNode, UntypedExpressionStatementNode, UntypedFunctionReturnNode, UntypedValNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index 8462d937..bfb9c822 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -118,7 +118,7 @@ public class Parser {
 
             body.add(statement);
 
-            if (tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE) || statement instanceof UntypedFunctionReturnNode) {
+            if (tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE) || statement instanceof UntypedBlockReturnNode) {
                 tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);
                 return body;
             }
@@ -193,18 +193,13 @@ public class Parser {
 
         var expression = parseTopLevelExpression(tokens);
 
-        // TODO: if we're not at the top level of a function (e.g. inside an
-        // if), then we're not going to return (from the function). What's
-        // better terminology for this? Java uses yield, but that has other
-        // meanings in other languages. Perhaps explicitly distinguish
-        // "block return" and "function return"?
         var isReturn = !tokens.trySkip(TokenType.SYMBOL_SEMICOLON);
 
         var end = tokens.position();
         var source = start.to(end);
 
         return isReturn
-            ? new UntypedFunctionReturnNode(expression, source)
+            ? new UntypedBlockReturnNode(expression, source)
             : new UntypedExpressionStatementNode(expression, source);
     }
 
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);
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 074aa7a6..43713d8a 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -69,6 +69,9 @@ public class TypeChecker {
 
     private static List<UntypedTargetNode.Variable> targetVariables(UntypedFunctionStatementNode statement) {
         return switch (statement) {
+            case UntypedBlockReturnNode blockReturn ->
+                List.of();
+
             case UntypedExpressionStatementNode expressionStatement ->
                 List.of();
 
@@ -199,6 +202,18 @@ public class TypeChecker {
         throw new UnsupportedBinaryOperationError(operator, left, right, source);
     }
 
+    private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckBlockReturn(
+        UntypedBlockReturnNode untyped,
+        TypeCheckerContext context
+    ) {
+        return typeCheckExpression(untyped.expression(), context).map(typedExpression -> {
+            return new TypedBlockReturnNode(
+                typedExpression,
+                untyped.source()
+            );
+        });
+    }
+
     public static TypeCheckerResult<TypedExpressionNode> typeCheckBoolLiteral(
         UntypedBoolLiteralNode untyped,
         TypeCheckerContext context
@@ -353,27 +368,31 @@ public class TypeChecker {
         bodyContext = bodyResult.context();
 
         var lastStatement = Lists.last(typedBody);
-        var actualReturnType = lastStatement
-            .flatMap(statement -> {
-                return switch (statement) {
-                    case TypedExpressionStatementNode typedExpressionStatement ->
-                        Optional.empty();
-
-                    case TypedFunctionReturnNode typedReturn ->
-                        Optional.of(typedReturn.expression().type());
-
-                    case TypedValNode typedVal ->
-                        Optional.empty();
-                };
-            })
-            .orElse(Types.UNIT);
-
-        if (!isSubType(actualReturnType, typedReturnType.value(), context)) {
-            throw new UnexpectedTypeError(
-                TypeSets.type(typedReturnType.value()),
-                actualReturnType,
-                untyped.source()
-            );
+        var lastStatementBehaviour = lastStatement
+            .map(statement -> statement.behaviour())
+            .orElse(TypedFunctionStatementBehaviour.NONE);
+
+        switch (lastStatementBehaviour) {
+            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                if (!isSubType(blockReturn.type(), typedReturnType.value(), context)) {
+                    throw new UnexpectedTypeError(
+                        TypeSets.type(typedReturnType.value()),
+                        blockReturn.type(),
+                        untyped.source()
+                    );
+                }
+            }
+            case TypedFunctionStatementBehaviour.None none -> {
+                if (!isSubType(Types.UNIT, typedReturnType.value(), context)) {
+                    throw new UnexpectedTypeError(
+                        TypeSets.type(typedReturnType.value()),
+                        Types.UNIT,
+                        untyped.source()
+                    );
+                }
+            }
+            case TypedFunctionStatementBehaviour.Terminate terminate -> {
+            }
         }
 
         var typed = new TypedFunctionDefinitionNode(
@@ -416,6 +435,9 @@ public class TypeChecker {
         TypeCheckerContext context
     ) {
         return switch (untyped) {
+            case UntypedBlockReturnNode blockReturn ->
+                typeCheckBlockReturn(blockReturn, context);
+
             case UntypedExpressionStatementNode expressionStatement ->
                 typeCheckExpressionStatement(expressionStatement, context);