Parse implicit return

1 files changed, 42 insertions, 1 deletions

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()));
+    }
 }