Throw error if function reaches end with inferred non-unit type

3 files changed, 59 insertions, 1 deletions

diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNodeMatcher.java
index d392ec4d..f31f1e0b 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNodeMatcher.java
@@ -3,8 +3,10 @@ package org.shedlang.compiler.ast.typed;
 import org.pcollections.Empty;
 import org.pcollections.PVector;
 import org.shedlang.compiler.testing.MatcherBuilder;
+import org.shedlang.compiler.types.Type;
 import org.zwobble.precisely.Matcher;
 
+import static org.zwobble.precisely.Matchers.equalTo;
 import static org.zwobble.precisely.Matchers.has;
 
 public class TypedFunctionBodyNodeMatcher extends MatcherBuilder<
@@ -43,4 +45,8 @@ public class TypedFunctionBodyNodeMatcher extends MatcherBuilder<
     ) {
         return with(has("freeVariables", x -> x.freeVariables(), freeVariables));
     }
+
+    public TypedFunctionBodyNodeMatcher withType(Type type) {
+        return with(has("type", x -> x.type(), equalTo(type)));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java
index be83f50d..5bc73325 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java
@@ -217,6 +217,57 @@ public class TypeCheckerFunctionBodyTests {
     }
 
     @Test
+    public void givenNoExplicitReturnTypeWhenNoExplicitReturnsThenReturnTypeIsUnit() {
+        var untyped = Untyped.block(List.of(
+            Untyped.expressionStatement(Untyped.stringLiteral(""))
+        ));
+        var signature = TypedFunctionSignatureNode.builder()
+            .withoutReturnType()
+            .build();
+        var context = TypeCheckerContext.EMPTY;
+
+        var typed = typeCheckFunctionBody(untyped, signature, context);
+
+        assertThat(
+            typed.value(),
+            isTypedFunctionBodyNode()
+                .withType(Types.UNIT)
+        );
+    }
+
+    @Test
+    public void givenNoExplicitReturnTypeWhenThereAreExplicitReturnsThenAllBranchesMustReturn() {
+        var untyped = Untyped.block(List.of(
+            Untyped.expressionStatement(Untyped.if_(
+                List.of(
+                    Untyped.conditionalBranch(
+                        Untyped.reference("c"),
+                        Untyped.block(List.of(
+                            Untyped.functionReturn(Untyped.stringLiteral(""))
+                        ))
+                    )
+                ),
+                Untyped.block(List.of())
+            ))
+        ));
+        var signature = TypedFunctionSignatureNode.builder()
+            .withoutReturnType()
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map(
+                Map.entry("c", Types.BOOL)
+            ));
+
+        var error = assertThrows(
+            UnexpectedTypeError.class,
+            () -> typeCheckFunctionBody(untyped, signature, context)
+        );
+
+        assertThat(error.expected(), equalTo(TypeSets.type(Types.STRING)));
+        assertThat(error.actual(), equalTo(Types.UNIT));
+    }
+
+    @Test
     public void freeVariablesIncludeVariablesReferencedInFunctionButDefinedOutsideOfFunction() {
         var untyped = Untyped.block(List.of(
             Untyped.expressionStatement(Untyped.reference("y"))
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java
index 1d858123..80627018 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java
@@ -40,9 +40,10 @@ public class TypeCheckerFunctionExpressionTests {
     }
 
     @Test
-    public void whenFunctionHasNoExplicitReturnsThenReturnTypeIsUnit() {
+    public void returnTypeCanBeInferred() {
         var untyped = UntypedFunctionExpressionNode.builder()
             .withoutReturnType()
+            .addStatement(Untyped.expressionStatement(Untyped.stringLiteral("")))
             .build();
         var context = TypeCheckerContext.EMPTY;