Support type-checking return statements with no expression

1 files changed, 19 insertions, 1 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index d00cc352..6dd1a3f3 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -435,6 +435,19 @@ public class TypeChecker {
         };
     }
 
+    private static TypeCheckerResult<? extends TypedExpressionNode> typeCheckExpressionOptional(
+        Optional<UntypedExpressionNode> untyped,
+        Optional<Type> hint,
+        Source source,
+        TypeCheckerContext context
+    ) {
+        if (untyped.isPresent()) {
+            return typeCheckExpression(untyped.get(), hint, context);
+        } else {
+            return context.toResult(new TypedUnitLiteralNode(source));
+        }
+    }
+
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckExpressionStatement(
         UntypedExpressionStatementNode untyped,
         TypeCheckerContext context
@@ -768,7 +781,12 @@ public class TypeChecker {
         TypeCheckerContext context
     ) {
         // TODO: pass return type as hint if available
-        var expressionResult = typeCheckExpression(untyped.expression().orElseThrow(), Optional.empty(), context);
+        var expressionResult = typeCheckExpressionOptional(
+            untyped.expression(),
+            Optional.empty(),
+            untyped.source(),
+            context
+        );
         var expression = expressionResult.value();
         context = expressionResult.context();