Handle block returns properly

11 files changed, 130 insertions, 21 deletions

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);