Support type-checking return statements with no expression

3 files changed, 24 insertions, 1 deletions

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