Parse return statements without expression

3 files changed, 7 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
index 94eb7a65..9809ce1f 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
@@ -2,8 +2,10 @@ package org.shedlang.compiler.ast.untyped;
 
 import org.shedlang.compiler.sources.Source;
 
+import java.util.Optional;
+
 public record UntypedFunctionReturnNode(
-    UntypedExpressionNode expression,
+    Optional<UntypedExpressionNode> expression,
     Source source
 ) implements UntypedFunctionStatementNode {
     @Override
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index 329a5283..272c66c0 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -413,7 +413,9 @@ public class Parser {
 
         tokens.skip(TokenType.KEYWORD_RETURN);
 
-        var expression = parseTopLevelExpression(tokens);
+        var expression = tokens.isNext(TokenType.SYMBOL_SEMICOLON)
+            ? Optional.<UntypedExpressionNode>empty()
+            : Optional.of(parseTopLevelExpression(tokens));
 
         tokens.skip(TokenType.SYMBOL_SEMICOLON);
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 00797708..d00cc352 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -768,7 +768,7 @@ public class TypeChecker {
         TypeCheckerContext context
     ) {
         // TODO: pass return type as hint if available
-        var expressionResult = typeCheckExpression(untyped.expression(), Optional.empty(), context);
+        var expressionResult = typeCheckExpression(untyped.expression().orElseThrow(), Optional.empty(), context);
         var expression = expressionResult.value();
         context = expressionResult.context();