Use hint for block returns

2 files changed, 70 insertions, 12 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
index be31d0a7..952a100a 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
@@ -5,10 +5,14 @@ import org.shedlang.compiler.types.Type;
 import org.shedlang.compiler.types.UserDefinedEffect;
 import org.shedlang.compiler.util.P;
 
+import java.util.Optional;
+
 public sealed interface ReturnTypeState {
     ReturnTypeState NOT_IN_FUNCTION = new NotInFunction();
     ReturnTypeState IMPLICIT_UNKNOWN = new ImplicitUnknown();
 
+    Optional<Type> returnType();
+
     static ReturnTypeState explicit(Type type) {
         return new Explicit(type);
     }
@@ -24,10 +28,33 @@ public sealed interface ReturnTypeState {
         return new EffectHandler(handledEffect, resumeType, P.vector());
     }
 
-    record NotInFunction() implements ReturnTypeState {}
-    record Explicit(Type type) implements ReturnTypeState {}
-    record ImplicitUnknown() implements ReturnTypeState {}
-    record ImplicitLowerBound(Type type) implements ReturnTypeState {}
+    record NotInFunction() implements ReturnTypeState {
+        @Override
+        public Optional<Type> returnType() {
+            return Optional.empty();
+        }
+    }
+
+    record Explicit(Type type) implements ReturnTypeState {
+        @Override
+        public Optional<Type> returnType() {
+            return Optional.of(type);
+        }
+    }
+
+    record ImplicitUnknown() implements ReturnTypeState {
+        @Override
+        public Optional<Type> returnType() {
+            return Optional.empty();
+        }
+    }
+
+    record ImplicitLowerBound(Type type) implements ReturnTypeState {
+        @Override
+        public Optional<Type> returnType() {
+            return Optional.of(type);
+        }
+    }
 
     record EffectHandler(
         UserDefinedEffect handledEffect,
@@ -37,5 +64,10 @@ public sealed interface ReturnTypeState {
         public ReturnTypeState addBreakType(Type breakType) {
             return new EffectHandler(handledEffect, resumeType, breakTypes.plus(breakType));
         }
+
+        @Override
+        public Optional<Type> returnType() {
+            return Optional.empty();
+        }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 69e326e9..ed658460 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -247,11 +247,16 @@ public class TypeChecker {
     private static TypeCheckerResult<TypedBlockNode> typeCheckBlock(
         UntypedBlockNode untyped,
         List<BoundVariable> declaredVariables,
+        Optional<Type> hint,
         TypeCheckerContext context
     ) {
         var typedStatementsResult = context.mapWithContext(
             untyped.statements(),
-            TypeChecker::typeCheckFunctionStatement
+            (untypedStatement, c) -> typeCheckFunctionStatement(
+                untypedStatement,
+                hint,
+                c
+            )
         );
         var typedStatements = typedStatementsResult.value();
         context = typedStatementsResult.context();
@@ -277,7 +282,12 @@ public class TypeChecker {
 
         context = context.enterBlock(bodyEnvironment.build());
 
-        var typedBodyResult = typeCheckBlock(untyped, bodyEnvironment.toBoundVariables(), context);
+        var typedBodyResult = typeCheckBlock(
+            untyped,
+            bodyEnvironment.toBoundVariables(),
+            Optional.empty(),
+            context
+        );
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
@@ -288,10 +298,10 @@ public class TypeChecker {
 
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckBlockReturn(
         UntypedBlockReturnNode untyped,
+        Optional<Type> hint,
         TypeCheckerContext context
     ) {
-        // TODO: pass return type as hint if available and applicable
-        return typeCheckExpression(untyped.expression(), Optional.empty(), context).map(typedExpression -> {
+        return typeCheckExpression(untyped.expression(), hint, context).map(typedExpression -> {
             return new TypedBlockReturnNode(
                 typedExpression,
                 untyped.source()
@@ -742,7 +752,12 @@ public class TypeChecker {
             signature.returnTypeState()
         );
 
-        var bodyResult = typeCheckBlock(untyped, bodyEnvironment.toBoundVariables(), context);
+        var bodyResult = typeCheckBlock(
+            untyped,
+            bodyEnvironment.toBoundVariables(),
+            signature.returnTypeState().returnType(),
+            context
+        );
         var typedBody = bodyResult.value();
         context = bodyResult.context();
 
@@ -1174,6 +1189,7 @@ public class TypeChecker {
 
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
         UntypedFunctionStatementNode untyped,
+        Optional<Type> blockReturnHint,
         TypeCheckerContext context
     ) {
         return switch (untyped) {
@@ -1181,7 +1197,7 @@ public class TypeChecker {
                 typeCheckAssign(assign, context);
 
             case UntypedBlockReturnNode blockReturn ->
-                typeCheckBlockReturn(blockReturn, context);
+                typeCheckBlockReturn(blockReturn, blockReturnHint, context);
 
             case UntypedBreakNode break_ ->
                 typeCheckBreak(break_, context);
@@ -1952,7 +1968,12 @@ public class TypeChecker {
 
         context = context.enterTryBody(bodyEnvironment.build(), typedHandledEffect.value());
 
-        var typedBodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.toBoundVariables(), context);
+        var typedBodyResult = typeCheckBlock(
+            untyped.body(),
+            bodyEnvironment.toBoundVariables(),
+            Optional.empty(),
+            context
+        );
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
@@ -2246,7 +2267,12 @@ public class TypeChecker {
             context = context.updateType(variableName.get(), variant, untyped.source());
         }
 
-        var typedBodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.toBoundVariables(), context);
+        var typedBodyResult = typeCheckBlock(
+            untyped.body(),
+            bodyEnvironment.toBoundVariables(),
+            Optional.empty(),
+            context
+        );
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();