Support type-checking return statements with no expression

4 files changed, 43 insertions, 2 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();
 
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index 0d60b280..a5a7456b 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -352,6 +352,10 @@ public class TypedNodeMatchers {
         return isTypedTypeLevelReferenceNode(name, equalTo(value));
     }
 
+    public static Matcher<TypedExpressionNode> isTypedUnitNode() {
+        return instanceOf(TypedUnitLiteralNode.class);
+    }
+
     public static Matcher<TypedNode> isTypedValNode(
         String targetName,
         Matcher<? super TypedExpressionNode> expression
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
index 5612ae3c..031dd362 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
@@ -94,7 +94,13 @@ public class Untyped {
     }
 
     public static UntypedFunctionReturnNode functionReturn(UntypedExpressionNode expression) {
-        return new UntypedFunctionReturnNode(Optional.of(expression), Sources.NULL);
+        return functionReturn(Optional.of(expression));
+    }
+
+    public static UntypedFunctionReturnNode functionReturn(
+        Optional<UntypedExpressionNode> expression
+    ) {
+        return new UntypedFunctionReturnNode(expression, Sources.NULL);
     }
 
     public static UntypedIfNode if_(
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
index a8ab9cb1..32908c34 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
@@ -10,6 +10,8 @@ import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
 
+import java.util.Optional;
+
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
 import static org.zwobble.precisely.AssertThat.assertThat;
@@ -28,6 +30,17 @@ public class TypeCheckerFunctionReturnTests {
     }
 
     @Test
+    public void whenExpressionIsAbsentThenUnitIsReturned() {
+        var untyped = Untyped.functionReturn(Optional.empty());
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, Types.UNIT);
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedFunctionReturnNode(isTypedUnitNode()));
+    }
+
+    @Test
     public void whenNotInFunctionThenReturnIsForbidden() {
         var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY