From 1e6ea9f8b84625c46a09b16fc7f68d04205465bd Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Tue, 7 Nov 2023 17:01:06 +0000
Subject: Use type hint when typing expression of return statements

---
 .../TypeCheckerFunctionReturnTests.java            | 24 ++++++++++++++++++++++
 1 file changed, 24 insertions(+)

(limited to 'src/test/java')

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)
+        ));
+    }
 }
-- 
cgit v1.2.3