Handle block returns properly

22 files changed, 254 insertions, 53 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);
 
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
index ab575632..d52f67a5 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
@@ -13,6 +13,10 @@ public class Typed {
     private Typed() {
     }
 
+    public static TypedBlockReturnNode blockReturn(TypedExpressionNode expression) {
+        return new TypedBlockReturnNode(expression, Sources.NULL);
+    }
+
     public static TypedBoolLiteralNode boolLiteral(boolean value) {
         return new TypedBoolLiteralNode(value, Sources.NULL);
     }
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index befeb326..56f73a12 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -29,6 +29,15 @@ public class TypedNodeMatchers {
         );
     }
 
+    public static Matcher<TypedFunctionStatementNode> isTypedBlockReturnNode(
+        Matcher<TypedExpressionNode> expression
+    ) {
+        return instanceOf(
+            TypedBlockReturnNode.class,
+            has("expression", x -> x.expression(), expression)
+        );
+    }
+
     public static Matcher<TypedExpressionNode> isTypedBoolLiteralNode(boolean value) {
         return instanceOf(
             TypedBoolLiteralNode.class,
@@ -53,6 +62,15 @@ public class TypedNodeMatchers {
         return TypedFunctionDefinitionNodeMatcher.empty();
     }
 
+    public static Matcher<TypedFunctionStatementNode> isTypedFunctionReturnNode(
+        Matcher<TypedExpressionNode> expression
+    ) {
+        return instanceOf(
+            TypedFunctionReturnNode.class,
+            has("expression", x -> x.expression(), expression)
+        );
+    }
+
     public static Matcher<TypedExpressionNode> isTypedIntAddNode(
         Matcher<TypedExpressionNode> left,
         Matcher<TypedExpressionNode> right
@@ -144,15 +162,6 @@ public class TypedNodeMatchers {
         );
     }
 
-    public static Matcher<TypedFunctionStatementNode> isTypedReturnNode(
-        Matcher<TypedExpressionNode> expression
-    ) {
-        return instanceOf(
-            TypedFunctionReturnNode.class,
-            has("expression", x -> x.expression(), expression)
-        );
-    }
-
     public static Matcher<TypedExpressionNode> isTypedStringAddNode(
         Matcher<TypedExpressionNode> left,
         Matcher<TypedExpressionNode> right
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
index db7037ce..24ca7f87 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
@@ -22,6 +22,10 @@ public class Untyped {
         );
     }
 
+    public static UntypedBlockReturnNode blockReturn(UntypedExpressionNode expression) {
+        return new UntypedBlockReturnNode(expression, Sources.NULL);
+    }
+
     public static UntypedBoolLiteralNode boolLiteral(boolean value) {
         return new UntypedBoolLiteralNode(value, Sources.NULL);
     }
@@ -62,7 +66,7 @@ public class Untyped {
         return new UntypedExpressionStatementNode(expression, Sources.NULL);
     }
 
-    public static UntypedFunctionReturnNode functionReturn_(UntypedExpressionNode expression) {
+    public static UntypedFunctionReturnNode functionReturn(UntypedExpressionNode expression) {
         return new UntypedFunctionReturnNode(expression, Sources.NULL);
     }
 
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
index 54edfc9d..dce29e55 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
@@ -34,6 +34,15 @@ public class UntypedNodeMatchers {
         );
     }
 
+    public static Matcher<UntypedFunctionStatementNode> isUntypedBlockReturnNode(
+        Matcher<UntypedExpressionNode> expression
+    ) {
+        return instanceOf(
+            UntypedBlockReturnNode.class,
+            has("expression", x -> x.expression(), expression)
+        );
+    }
+
     public static UntypedCallNodeMatcher isUntypedCallNode() {
         return UntypedCallNodeMatcher.empty();
     }
diff --git a/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java
index 88bde61d..279309f6 100644
--- a/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java
@@ -205,7 +205,7 @@ public class ParserFunctionDefinitionTests {
     }
 
     @Test
-    public void canParseFunctionWithBodyWithImplicitReturn() {
+    public void canParseFunctionWithBodyWithBlockReturn() {
         var source = """
             fun f() -> Unit {
                 a;
@@ -221,13 +221,13 @@ public class ParserFunctionDefinitionTests {
                 .withBody(isSequence(
                     isUntypedExpressionStatementNode(isUntypedReferenceNode("a")),
                     isUntypedExpressionStatementNode(isUntypedReferenceNode("b")),
-                    isUntypedFunctionReturnNode(isUntypedReferenceNode("c"))
+                    isUntypedBlockReturnNode(isUntypedReferenceNode("c"))
                 ))
         );
     }
 
     @Test
-    public void implicitReturnEndsFunction() {
+    public void blockReturnEndsFunction() {
         var source = """
             fun f() -> Unit {
                 a;
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerBlockReturnTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerBlockReturnTests.java
new file mode 100644
index 00000000..7ff57a10
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerBlockReturnTests.java
@@ -0,0 +1,21 @@
+package org.shedlang.compiler.stackir;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.typed.Typed;
+
+import static org.shedlang.compiler.stackir.instructions.InstructionMatchers.isPushStringInstruction;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.isSequence;
+
+public class StackCompilerBlockReturnTests {
+    @Test
+    public void expressionIsEvaluated() {
+        var node = Typed.blockReturn(Typed.stringLiteral("hello"));
+
+        var instructions = StackCompiler.compileFunctionStatement(node);
+
+        assertThat(instructions, isSequence(
+            isPushStringInstruction("hello")
+        ));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
index 28974ea9..051b0fb6 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionDefinitionTests.java
@@ -46,4 +46,28 @@ public class StackCompilerFunctionDefinitionTests {
             isLocalStoreInstruction("f")
         ));
     }
+
+    @Test
+    public void whenLastStatementIsBlockReturnThenThatValueIsReturned() {
+        var node = TypedFunctionDefinitionNode.builder()
+            .withName("f")
+            .addPositionalParam("x", Typed.typeLevelReference("String", Types.STRING))
+            .addNamedParam("y", Typed.typeLevelReference("Int", Types.INT))
+            .addFreeVariable("z")
+            .addStatement(Typed.blockReturn(Typed.stringLiteral("hello")))
+            .build();
+
+        var instructions = StackCompiler.compileNamespaceStatement(node);
+
+        assertThat(instructions, isSequence(
+            instanceOf(
+                DefineFunctionInstruction.class,
+                has("body", x -> x.body(), isSequence(
+                    isPushStringInstruction("hello"),
+                    isReturnInstruction()
+                ))
+            ),
+            isLocalStoreInstruction("f")
+        ));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerReturnTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionReturnTests.java
index dcb1af47..0baf9acb 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerFunctionReturnTests.java
@@ -7,7 +7,7 @@ import static org.shedlang.compiler.stackir.instructions.InstructionMatchers.*;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.isSequence;
 
-public class StackCompilerReturnTests {
+public class StackCompilerFunctionReturnTests {
     @Test
     public void expressionIsEvaluatedAndTheResultDiscarded() {
         var node = Typed.functionReturn(Typed.stringLiteral("hello"));
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java
new file mode 100644
index 00000000..bf4a948f
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java
@@ -0,0 +1,24 @@
+package org.shedlang.compiler.typechecker;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.untyped.Untyped;
+import org.shedlang.compiler.types.Effects;
+import org.shedlang.compiler.types.Types;
+import org.shedlang.compiler.util.P;
+
+import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedBlockReturnNode;
+import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedStringLiteralNode;
+import static org.zwobble.precisely.AssertThat.assertThat;
+
+public class TypeCheckerBlockReturnTests {
+    @Test
+    public void expressionIsTypeChecked() {
+        var untyped = Untyped.blockReturn(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, Types.STRING);
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedBlockReturnNode(isTypedStringLiteralNode("hello")));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index 2ff53f74..4c640f30 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -205,7 +205,7 @@ public class TypeCheckerFunctionDefinitionTests {
     public void bodyHasReturnTypeAsDefinedInSignature() {
         var untyped = UntypedFunctionDefinitionNode.builder()
             .withReturnType(Untyped.typeLevelReference("String"))
-            .addStatement(Untyped.functionReturn_(Untyped.intLiteral(42)))
+            .addStatement(Untyped.functionReturn(Untyped.intLiteral(42)))
             .build();
         var context = TypeCheckerContext.EMPTY
             .enterNamespaceWithDefinedVariables(P.map(
@@ -254,6 +254,20 @@ public class TypeCheckerFunctionDefinitionTests {
     }
 
     @Test
+    public void givenReturnTypeIsNotUnitWhenBodyHasBlockReturnOfCorrectTypeThenFunctionTypeChecks() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .withReturnType(Untyped.typeLevelReference("String"))
+            .addStatement(Untyped.blockReturn(Untyped.stringLiteral("hello")))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespaceWithDefinedVariables(P.map(
+                Map.entry("String", Types.metaType(Types.STRING))
+            ));
+
+        TypeChecker.typeCheckNamespaceStatement(untyped, context);
+    }
+
+    @Test
     public void freeVariablesIncludeVariablesReferencedInFunctionButDefinedOutsideOfFunction() {
         var untyped = UntypedFunctionDefinitionNode.builder()
             .withReturnType(Untyped.typeLevelReference("Unit"))
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
index bea2d4d9..2e4b4206 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
@@ -17,18 +17,18 @@ import static org.zwobble.precisely.Matchers.equalTo;
 public class TypeCheckerFunctionReturnTests {
     @Test
     public void expressionIsTypeChecked() {
-        var untyped = Untyped.functionReturn_(Untyped.stringLiteral("hello"));
+        var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.STRING);
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
-        assertThat(typed.value(), isTypedReturnNode(isTypedStringLiteralNode("hello")));
+        assertThat(typed.value(), isTypedFunctionReturnNode(isTypedStringLiteralNode("hello")));
     }
 
     @Test
     public void whenNotInFunctionThenReturnIsForbidden() {
-        var untyped = Untyped.functionReturn_(Untyped.stringLiteral("hello"));
+        var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
             .enterNamespace(P.map());
 
@@ -40,18 +40,18 @@ public class TypeCheckerFunctionReturnTests {
 
     @Test
     public void canReturnSubtypeOfReturnType() {
-        var untyped = Untyped.functionReturn_(Untyped.stringLiteral("hello"));
+        var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.TOP);
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
-        assertThat(typed.value(), isTypedReturnNode(isTypedStringLiteralNode("hello")));
+        assertThat(typed.value(), isTypedFunctionReturnNode(isTypedStringLiteralNode("hello")));
     }
 
     @Test
     public void cannotReturnNonSubtypeOfReturnType() {
-        var untyped = Untyped.functionReturn_(Untyped.stringLiteral("hello"));
+        var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.INT);