Parse implicit return

2 files changed, 51 insertions, 1 deletions

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 ad6c8a49..1f74008d 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
@@ -79,6 +79,15 @@ public class UntypedNodeMatchers {
         );
     }
 
+    public static Matcher<UntypedFunctionStatementNode> isUntypedReturnNode(
+        Matcher<UntypedExpressionNode> expression
+    ) {
+        return instanceOf(
+            UntypedReturnNode.class,
+            has("expression", x -> x.expression(), expression)
+        );
+    }
+
     public static Matcher<UntypedExpressionNode> isUntypedStringLiteralNode(String value) {
         return instanceOf(
             UntypedStringLiteralNode.class,
diff --git a/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java
index c9961426..0f96b59b 100644
--- a/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/parser/ParserFunctionDefinitionTests.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.parser;
 
 import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.tokeniser.UnexpectedTokenException;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.untyped.UntypedNodeMatchers.*;
@@ -182,7 +183,7 @@ public class ParserFunctionDefinitionTests {
     }
 
     @Test
-    public void canParseFunctionWithBody() {
+    public void canParseFunctionWithBodyWithoutReturn() {
         var source = """
             fun f() -> Unit {
                 a;
@@ -202,4 +203,44 @@ public class ParserFunctionDefinitionTests {
                 ))
         );
     }
+
+    @Test
+    public void canParseFunctionWithBodyWithImplicitReturn() {
+        var source = """
+            fun f() -> Unit {
+                a;
+                b;
+                c
+            }""";
+
+        var node = parseString(source, Parser::parseNamespaceStatement);
+
+        assertThat(
+            node,
+            isUntypedFunctionDefinitionNode()
+                .withBody(isSequence(
+                    isUntypedExpressionStatementNode(isUntypedReferenceNode("a")),
+                    isUntypedExpressionStatementNode(isUntypedReferenceNode("b")),
+                    isUntypedReturnNode(isUntypedReferenceNode("c"))
+                ))
+        );
+    }
+
+    @Test
+    public void implicitReturnEndsFunction() {
+        var source = """
+            fun f() -> Unit {
+                a;
+                b
+                c
+            }""";
+
+        var error = assertThrows(
+            UnexpectedTokenException.class,
+            () -> parseString(source, Parser::parseNamespaceStatement)
+        );
+
+        assertThat(error.getActual().value(), equalTo("c"));
+        assertThat(error.getExpected(), equalTo(TokenType.SYMBOL_BRACE_CLOSE.toString()));
+    }
 }