Use type hint when typing expression of return statements

1 files changed, 24 insertions, 0 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
index 650d1531..f0c2daa7 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
@@ -2,11 +2,13 @@ package org.shedlang.compiler.typechecker;
 
 import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.ast.untyped.Untyped;
+import org.shedlang.compiler.ast.untyped.UntypedFunctionExpressionNode;
 import org.shedlang.compiler.typechecker.errors.CannotReturnFromEffectHandlerFieldError;
 import org.shedlang.compiler.typechecker.errors.CannotReturnOutsideOfFunctionError;
 import org.shedlang.compiler.typechecker.errors.UnexpectedTypeError;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.Effects;
+import org.shedlang.compiler.types.FunctionType;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
@@ -90,4 +92,26 @@ public class TypeCheckerFunctionReturnTests {
         assertThat(error.expected(), equalTo(TypeSets.type(Types.INT)));
         assertThat(error.actual(), equalTo(Types.STRING));
     }
+
+    @Test
+    public void returnTypeIsUsedAsTypeHintForExpression() {
+        var untyped = Untyped.functionReturn(
+            UntypedFunctionExpressionNode.builder()
+                .addPositionalParam("x", Optional.empty())
+                .withoutReturnType()
+                .build()
+        );
+        var returnType = FunctionType.builder()
+            .addPositionalParam(Types.STRING)
+            .withReturnType(Types.UNIT)
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, returnType);
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedFunctionReturnNode(
+            isTypedFunctionExpressionNode().withType(returnType)
+        ));
+    }
 }