Use type hint when typing expression of return statements

1 files changed, 5 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 756e4584..3e9f4a46 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -804,10 +804,13 @@ public class TypeChecker {
         UntypedFunctionReturnNode untyped,
         TypeCheckerContext context
     ) {
-        // TODO: pass return type as hint if available
+        var expressionTypeHint = context.returnTypeState() instanceof ReturnTypeState.Explicit returnTypeExplicit
+            ? Optional.of(returnTypeExplicit.type())
+            : Optional.<Type>empty();
+
         var expressionResult = typeCheckExpressionOptional(
             untyped.expression(),
-            Optional.empty(),
+            expressionTypeHint,
             untyped.source(),
             context
         );