Add explicit function return statement

9 files changed, 59 insertions, 39 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedReturnNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionReturnNode.java
index ad6d1f08..883b5649 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedReturnNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionReturnNode.java
@@ -2,7 +2,7 @@ package org.shedlang.compiler.ast.typed;
 
 import org.shedlang.compiler.sources.Source;
 
-public record TypedReturnNode(
+public record TypedFunctionReturnNode(
     TypedExpressionNode expression,
     Source source
 ) implements TypedFunctionStatementNode {
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 ac46cde4..bbc0fc6f 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,4 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedExpressionStatementNode, TypedReturnNode, TypedValNode {
+public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedExpressionStatementNode, TypedFunctionReturnNode, TypedValNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedReturnNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
index 36230db8..c3f5663e 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedReturnNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
@@ -2,7 +2,7 @@ package org.shedlang.compiler.ast.untyped;
 
 import org.shedlang.compiler.sources.Source;
 
-public record UntypedReturnNode(
+public record UntypedFunctionReturnNode(
     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 6ecff062..d970d6f3 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, UntypedReturnNode, UntypedValNode {
+public sealed interface UntypedFunctionStatementNode extends UntypedNode permits 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 c7945d5a..8462d937 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 UntypedReturnNode) {
+            if (tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE) || statement instanceof UntypedFunctionReturnNode) {
                 tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);
                 return body;
             }
@@ -204,7 +204,7 @@ public class Parser {
         var source = start.to(end);
 
         return isReturn
-            ? new UntypedReturnNode(expression, source)
+            ? new UntypedFunctionReturnNode(expression, source)
             : new UntypedExpressionStatementNode(expression, source);
     }
 
@@ -243,8 +243,26 @@ public class Parser {
         );
     }
 
+    private static UntypedFunctionStatementNode parseFunctionReturn(TokenIterator<TokenType> tokens) {
+        var start = tokens.position();
+
+        tokens.skip(TokenType.KEYWORD_RETURN);
+
+        var expression = parseTopLevelExpression(tokens);
+
+        tokens.skip(TokenType.SYMBOL_SEMICOLON);
+
+        var end = tokens.position();
+        var source = start.to(end);
+
+        return new UntypedFunctionReturnNode(expression, source);
+    }
+
     static UntypedFunctionStatementNode parseFunctionStatement(TokenIterator<TokenType> tokens) {
         return switch (tokens.peek().tokenType()) {
+            case KEYWORD_RETURN ->
+                parseFunctionReturn(tokens);
+
             case KEYWORD_VAL ->
                 parseVal(tokens);
 
diff --git a/src/main/java/org/shedlang/compiler/parser/TokenType.java b/src/main/java/org/shedlang/compiler/parser/TokenType.java
index a7df78ff..337692de 100644
--- a/src/main/java/org/shedlang/compiler/parser/TokenType.java
+++ b/src/main/java/org/shedlang/compiler/parser/TokenType.java
@@ -11,6 +11,7 @@ public enum TokenType {
     KEYWORD_ELSE,
     KEYWORD_FUN,
     KEYWORD_IF,
+    KEYWORD_RETURN,
     KEYWORD_VAL,
 
     SYMBOL_BANG,
diff --git a/src/main/java/org/shedlang/compiler/parser/Tokeniser.java b/src/main/java/org/shedlang/compiler/parser/Tokeniser.java
index 3a2f3e11..99542eb0 100644
--- a/src/main/java/org/shedlang/compiler/parser/Tokeniser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Tokeniser.java
@@ -16,6 +16,7 @@ public class Tokeniser {
         Map.entry("else", TokenType.KEYWORD_ELSE),
         Map.entry("fun", TokenType.KEYWORD_FUN),
         Map.entry("if", TokenType.KEYWORD_IF),
+        Map.entry("return", TokenType.KEYWORD_RETURN),
         Map.entry("val", TokenType.KEYWORD_VAL)
     );
 
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 449360e1..d3b8801c 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -112,14 +112,19 @@ public class StackCompiler {
 
         return P.vector(defineFunctionInstruction, storeInstruction);
     }
+
+    private static PVector<Instruction> compileFunctionReturn(TypedFunctionReturnNode node) {
+        return compileExpression(node.expression())
+            .plus(new ReturnInstruction(node.source()));
+    }
     
     public static PVector<Instruction> compileFunctionStatement(TypedFunctionStatementNode node) {
         return switch (node) {
             case TypedExpressionStatementNode expressionStatement ->
                 compileExpressionStatement(expressionStatement);
 
-            case TypedReturnNode return_ ->
-                compileReturn(return_);
+            case TypedFunctionReturnNode functionReturn ->
+                compileFunctionReturn(functionReturn);
 
             case TypedValNode val ->
                 compileVal(val);
@@ -169,11 +174,6 @@ public class StackCompiler {
         return P.vector(new LocalLoadInstruction(node.name(), node.source()));
     }
 
-    private static PVector<Instruction> compileReturn(TypedReturnNode node) {
-        return compileExpression(node.expression())
-            .plus(new ReturnInstruction(node.source()));
-    }
-
     private static PVector<Instruction> compileStringLiteral(TypedStringLiteralNode node) {
         return P.vector(new PushStringInstruction(node.value(), node.source()));
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index b938c66e..074aa7a6 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -72,7 +72,7 @@ public class TypeChecker {
             case UntypedExpressionStatementNode expressionStatement ->
                 List.of();
 
-            case UntypedReturnNode return_ ->
+            case UntypedFunctionReturnNode functionReturn ->
                 List.of();
 
             case UntypedValNode val ->
@@ -359,7 +359,7 @@ public class TypeChecker {
                     case TypedExpressionStatementNode typedExpressionStatement ->
                         Optional.empty();
 
-                    case TypedReturnNode typedReturn ->
+                    case TypedFunctionReturnNode typedReturn ->
                         Optional.of(typedReturn.expression().type());
 
                     case TypedValNode typedVal ->
@@ -389,6 +389,28 @@ public class TypeChecker {
         return context.define(untyped.name(), typed.type()).toResult(typed);
     }
 
+    private static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionReturn(
+        UntypedFunctionReturnNode untyped,
+        TypeCheckerContext context
+    ) {
+        var expressionResult = typeCheckExpression(untyped.expression(), context);
+        var expression = expressionResult.value();
+        context = expressionResult.context();
+
+        if (context.returnType().isEmpty()) {
+            // This should be syntactically impossible, but since we only have
+            // one context type, we check for this possibility and throw a
+            // reasonable error.
+            throw new CannotReturnOutsideOfFunction(untyped.source());
+        }
+
+        checkExpressionType(expression, context.returnType().get(), context);
+
+        var typed = new TypedFunctionReturnNode(expression, untyped.source());
+
+        return context.toResult(typed);
+    }
+
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
         UntypedFunctionStatementNode untyped,
         TypeCheckerContext context
@@ -397,8 +419,8 @@ public class TypeChecker {
             case UntypedExpressionStatementNode expressionStatement ->
                 typeCheckExpressionStatement(expressionStatement, context);
 
-            case UntypedReturnNode return_ ->
-                typeCheckReturn(return_, context);
+            case UntypedFunctionReturnNode functionReturn ->
+                typeCheckFunctionReturn(functionReturn, context);
 
             case UntypedValNode val ->
                 typeCheckVal(val, context);
@@ -501,28 +523,6 @@ public class TypeChecker {
         });
     }
 
-    private static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckReturn(
-        UntypedReturnNode untyped,
-        TypeCheckerContext context
-    ) {
-        var expressionResult = typeCheckExpression(untyped.expression(), context);
-        var expression = expressionResult.value();
-        context = expressionResult.context();
-
-        if (context.returnType().isEmpty()) {
-            // This should be syntactically impossible, but since we only have
-            // one context type, we check for this possibility and throw a
-            // reasonable error.
-            throw new CannotReturnOutsideOfFunction(untyped.source());
-        }
-
-        checkExpressionType(expression, context.returnType().get(), context);
-
-        var typed = new TypedReturnNode(expression, untyped.source());
-
-        return context.toResult(typed);
-    }
-
     private static TypeCheckerResult<TypedStringLiteralNode> typeCheckStringLiteral(
         UntypedStringLiteralNode untyped,
         TypeCheckerContext context