Use hint for block returns

11 files changed, 83 insertions, 27 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerAssignTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerAssignTests.java
index 6c3cf09d..7bdf3c8b 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerAssignTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerAssignTests.java
@@ -173,6 +173,6 @@ public class TypeCheckerAssignTests {
 
         context = context.enterNamespace(NamespaceName.of("Example"), environment.build());
 
-        return TypeChecker.typeCheckFunctionStatement(untyped, context);
+        return TypeCheckerTesting.typeCheckFunctionStatement(untyped, context);
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java
index bf4a948f..8aa1351c 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBlockReturnTests.java
@@ -6,6 +6,8 @@ import org.shedlang.compiler.types.Effects;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
 
+import java.util.Optional;
+
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedBlockReturnNode;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedStringLiteralNode;
 import static org.zwobble.precisely.AssertThat.assertThat;
@@ -17,7 +19,11 @@ public class TypeCheckerBlockReturnTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.STRING);
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = TypeChecker.typeCheckFunctionStatement(
+            untyped,
+            Optional.empty(),
+            context
+        );
 
         assertThat(typed.value(), isTypedBlockReturnNode(isTypedStringLiteralNode("hello")));
     }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
index c355b0cb..8c8a1296 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
@@ -11,6 +11,7 @@ import org.shedlang.compiler.util.P;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.shedlang.compiler.typechecker.TypeCheckerTesting.typeCheckFunctionStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 
 public class TypeCheckerBreakTests {
@@ -23,7 +24,7 @@ public class TypeCheckerBreakTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.INT));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedBreakNode(
             EFFECT,
@@ -37,7 +38,7 @@ public class TypeCheckerBreakTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.UNIT));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedBreakNode(
             EFFECT,
@@ -53,7 +54,7 @@ public class TypeCheckerBreakTests {
 
         assertThrows(
             CannotBreakOutsideOfEffectHandlerMethodError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 
@@ -65,7 +66,7 @@ public class TypeCheckerBreakTests {
 
         assertThrows(
             CannotBreakOutsideOfEffectHandlerMethodError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java
index b7ffca1b..bba73a14 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java
@@ -5,6 +5,7 @@ import org.shedlang.compiler.ast.untyped.Untyped;
 
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedExpressionStatementNode;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedStringLiteralNode;
+import static org.shedlang.compiler.typechecker.TypeCheckerTesting.typeCheckFunctionStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 
 public class TypeCheckerExpressionStatementTests {
@@ -12,7 +13,7 @@ public class TypeCheckerExpressionStatementTests {
     public void expressionIsTypeChecked() {
         var untyped = Untyped.expressionStatement(Untyped.stringLiteral("hello"));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, TypeCheckerContext.EMPTY);
+        var typed = typeCheckFunctionStatement(untyped, TypeCheckerContext.EMPTY);
 
         assertThat(typed.value(), isTypedExpressionStatementNode(isTypedStringLiteralNode("hello")));
     }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java
index b2dc700e..bcf5e389 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionBodyTests.java
@@ -4,6 +4,7 @@ import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.ast.typed.BoundVariable;
 import org.shedlang.compiler.ast.untyped.Untyped;
 import org.shedlang.compiler.ast.untyped.UntypedCallNode;
+import org.shedlang.compiler.ast.untyped.UntypedFunctionExpressionNode;
 import org.shedlang.compiler.typechecker.errors.CannotCaptureVarError;
 import org.shedlang.compiler.typechecker.errors.DuplicateVariableNameError;
 import org.shedlang.compiler.typechecker.errors.UnexpectedTypeError;
@@ -16,6 +17,7 @@ import org.shedlang.compiler.util.P;
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
@@ -218,6 +220,36 @@ public class TypeCheckerFunctionBodyTests {
     }
 
     @Test
+    public void returnTypeIsUsedAsTypeHintForBlockReturn() {
+        var untyped = Untyped.block(List.of(
+            Untyped.blockReturn(
+                UntypedFunctionExpressionNode.builder()
+                    .addPositionalParam("x", Optional.empty())
+                    .withoutReturnType()
+                    .build()
+            )
+        ));
+        var returnType = FunctionType.builder()
+            .addPositionalParam(Types.STRING)
+            .withReturnType(Types.UNIT)
+            .build();
+        var signature = PartialFunctionSignature.builder()
+            .withReturnType(returnType)
+            .build();
+        var context = TypeCheckerContext.EMPTY;
+
+        var result = typeCheckFunctionBody(untyped, signature, context);
+
+        assertThat(result.value(), isTypedFunctionBodyNode()
+            .withBlock(isTypedBlockNode(isSequence(
+                isTypedBlockReturnNode(
+                    isTypedFunctionExpressionNode().withType(returnType)
+                )
+            )))
+        );
+    }
+
+    @Test
     public void givenNoExplicitReturnTypeWhenNoReturnStatementsThenReturnTypeIsUnit() {
         var untyped = Untyped.block(List.of(
             Untyped.expressionStatement(Untyped.stringLiteral(""))
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
index 6d8c75df..d0cc356f 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
@@ -14,6 +14,7 @@ import java.util.Optional;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.shedlang.compiler.typechecker.TypeCheckerTesting.typeCheckFunctionStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.equalTo;
 
@@ -24,7 +25,7 @@ public class TypeCheckerFunctionReturnTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.STRING);
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedFunctionReturnNode(isTypedStringLiteralNode("hello")));
     }
@@ -35,7 +36,7 @@ public class TypeCheckerFunctionReturnTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.UNIT);
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedFunctionReturnNode(isTypedUnitLiteralNode()));
     }
@@ -48,7 +49,7 @@ public class TypeCheckerFunctionReturnTests {
 
         assertThrows(
             CannotReturnOutsideOfFunctionError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 
@@ -65,7 +66,7 @@ public class TypeCheckerFunctionReturnTests {
 
         assertThrows(
             CannotReturnFromEffectHandlerMethodError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 
@@ -75,7 +76,7 @@ public class TypeCheckerFunctionReturnTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, Types.TOP);
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedFunctionReturnNode(isTypedStringLiteralNode("hello")));
     }
@@ -88,7 +89,7 @@ public class TypeCheckerFunctionReturnTests {
 
         var error = assertThrows(
             UnexpectedTypeError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
 
         assertThat(error.expected(), equalTo(TypeSets.type(Types.INT)));
@@ -110,7 +111,7 @@ public class TypeCheckerFunctionReturnTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, returnType);
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedFunctionReturnNode(
             isTypedFunctionExpressionNode().withType(returnType)
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
index 91ba8c74..21a8a4cb 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
@@ -13,6 +13,7 @@ import java.util.Optional;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.shedlang.compiler.typechecker.TypeCheckerTesting.typeCheckFunctionStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.equalTo;
 
@@ -26,7 +27,7 @@ public class TypeCheckerResumeTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.STRING));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedResumeNode(isTypedStringLiteralNode("hello")));
     }
@@ -37,7 +38,7 @@ public class TypeCheckerResumeTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.UNIT));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedResumeNode(isTypedUnitLiteralNode()));
     }
@@ -50,7 +51,7 @@ public class TypeCheckerResumeTests {
 
         assertThrows(
             CannotResumeOutsideOfEffectHandlerMethodError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 
@@ -62,7 +63,7 @@ public class TypeCheckerResumeTests {
 
         assertThrows(
             CannotResumeOutsideOfEffectHandlerMethodError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 
@@ -72,7 +73,7 @@ public class TypeCheckerResumeTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.TOP));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedResumeNode(isTypedStringLiteralNode("hello")));
     }
@@ -85,7 +86,7 @@ public class TypeCheckerResumeTests {
 
         var error = assertThrows(
             UnexpectedTypeError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
 
         assertThat(error.expected(), equalTo(TypeSets.type(Types.INT)));
@@ -107,7 +108,7 @@ public class TypeCheckerResumeTests {
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, resumeType));
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedResumeNode(
             isTypedFunctionExpressionNode().withType(resumeType)
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTesting.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTesting.java
index d98894f2..32f0e1ca 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTesting.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTesting.java
@@ -1,8 +1,10 @@
 package org.shedlang.compiler.typechecker;
 
 import org.shedlang.compiler.ast.typed.TypedExpressionNode;
+import org.shedlang.compiler.ast.typed.TypedFunctionStatementNode;
 import org.shedlang.compiler.ast.typed.TypedNamespaceStatementNode;
 import org.shedlang.compiler.ast.untyped.UntypedExpressionNode;
+import org.shedlang.compiler.ast.untyped.UntypedFunctionStatementNode;
 import org.shedlang.compiler.ast.untyped.UntypedNamespaceStatementNode;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.util.P;
@@ -25,6 +27,17 @@ public class TypeCheckerTesting {
         return TypeChecker.typeCheckStepsNamespaceStatement(untyped).typeCheckAllPhases(context);
     }
 
+    static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
+        UntypedFunctionStatementNode untyped,
+        TypeCheckerContext context
+    ) {
+        return TypeChecker.typeCheckFunctionStatement(
+            untyped,
+            Optional.empty(),
+            context
+        );
+    }
+
     static TypeCheckerResult<? extends TypedExpressionNode> typeCheckExpression(
         UntypedExpressionNode untyped,
         TypeCheckerContext context
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerValTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerValTests.java
index 3f2adf5d..39c5c1e5 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerValTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerValTests.java
@@ -162,6 +162,6 @@ public class TypeCheckerValTests {
 
         context = context.enterNamespace(NamespaceName.of("Example"), environment.build());
 
-        return TypeChecker.typeCheckFunctionStatement(untyped, context);
+        return TypeCheckerTesting.typeCheckFunctionStatement(untyped, context);
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerVarTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerVarTests.java
index 0443802f..7ac18c56 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerVarTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerVarTests.java
@@ -63,6 +63,6 @@ public class TypeCheckerVarTests {
 
         context = context.enterNamespace(NamespaceName.of("Example"), environment.build());
 
-        return TypeChecker.typeCheckFunctionStatement(untyped, context);
+        return TypeCheckerTesting.typeCheckFunctionStatement(untyped, context);
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java
index 5fc029b8..e84d59b8 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java
@@ -11,6 +11,7 @@ import java.util.List;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.shedlang.compiler.typechecker.TypeCheckerTesting.typeCheckFunctionStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.*;
 
@@ -23,7 +24,7 @@ public class TypeCheckerWhileTests {
         );
         var context = TypeCheckerContext.EMPTY;
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedWhileNode(
             isTypedBoolLiteralNode(true),
@@ -41,7 +42,7 @@ public class TypeCheckerWhileTests {
 
         var error = assertThrows(
             UnexpectedTypeError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
 
         assertThat(error.expected(), equalTo(TypeSets.type(Types.BOOL)));
@@ -58,7 +59,7 @@ public class TypeCheckerWhileTests {
         );
         var context = TypeCheckerContext.EMPTY;
 
-        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+        var typed = typeCheckFunctionStatement(untyped, context);
 
         assertThat(typed.value(), isTypedWhileNode(
             anything(),
@@ -80,7 +81,7 @@ public class TypeCheckerWhileTests {
 
         assertThrows(
             LoopBodyCannotYieldValueError.class,
-            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+            () -> typeCheckFunctionStatement(untyped, context)
         );
     }
 }