Type check resume statements properly

7 files changed, 181 insertions, 9 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
index 4edfe582..cfcd59ec 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
@@ -32,6 +32,15 @@ public record PartialFunctionSignature(
         );
     }
 
+    public PartialFunctionSignature withReturnTypeState(ReturnTypeState returnTypeState) {
+        return new PartialFunctionSignature(
+            this.positionalParams,
+            this.namedParams,
+            this.effect,
+            returnTypeState
+        );
+    }
+
     public static Builder builder() {
         return new Builder(
             P.vector(),
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 3e9f4a46..d04bcea4 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1290,12 +1290,42 @@ public class TypeChecker {
         UntypedResumeNode untyped,
         TypeCheckerContext context
     ) {
-        context = typeCheckFunctionReturnValue(Types.UNIT, untyped.source(), context);
+        switch (context.returnTypeState()) {
+            case ReturnTypeState.Explicit explicit -> {
+                throw new CannotResumeOutsideOfEffectHandlerFieldError(untyped.source());
+            }
+            case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
+                throw new CannotResumeOutsideOfEffectHandlerFieldError(untyped.source());
+            }
+            case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
+                throw new CannotResumeOutsideOfEffectHandlerFieldError(untyped.source());
+            }
+            case ReturnTypeState.NotInFunction notInFunction -> {
+                throw new CannotResumeOutsideOfEffectHandlerFieldError(untyped.source());
+            }
+            case ReturnTypeState.EffectHandler effectHandler -> {
+                var typedExpressionResult = typeCheckExpressionOptional(
+                    untyped.expression(),
+                    Optional.of(effectHandler.resumeType()),
+                    untyped.source(),
+                    context
+                );
+                var typedExpression = typedExpressionResult.value();
+                context = typedExpressionResult.context();
 
-        return context.toResult(new TypedResumeNode(
-            new TypedUnitLiteralNode(untyped.source()),
-            untyped.source()
-        ));
+                checkIsSubType(
+                    typedExpression.type(),
+                    effectHandler.resumeType(),
+                    typedExpression.source(),
+                    context
+                );
+
+                return context.toResult(new TypedResumeNode(
+                    typedExpression,
+                    untyped.source()
+                ));
+            }
+        }
     }
 
     private static TypeCheckerResult<TypedStringLiteralNode> typeCheckStringLiteral(
@@ -1697,9 +1727,12 @@ public class TypeChecker {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
         var partialSignature = inferFunctionPartialSignature(typedSignature, Optional.of(effectFieldType))
-            .withEffect(context.effect());
+            // TODO: check effect is none before overriding
+            .withEffect(context.effect())
+            // TODO: check return type is implicit unknown before overriding
+            .withReturnTypeState(ReturnTypeState.effectHandler(effectFieldType.returnType()));
+
 
-        // TODO: don't allow returns, allow resumes
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
             partialSignature,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/CannotResumeOutsideOfEffectHandlerFieldError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotResumeOutsideOfEffectHandlerFieldError.java
new file mode 100644
index 00000000..6d53c809
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotResumeOutsideOfEffectHandlerFieldError.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+
+public class CannotResumeOutsideOfEffectHandlerFieldError extends SourceError {
+    public CannotResumeOutsideOfEffectHandlerFieldError(Source source) {
+        super("Cannot return outside of effect handler field", source);
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
index adae966a..51b6c2b6 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
@@ -157,6 +157,10 @@ public class Untyped {
         return new UntypedResumeNode(Optional.empty(), Sources.NULL);
     }
 
+    public static UntypedResumeNode resume(UntypedExpressionNode expression) {
+        return new UntypedResumeNode(Optional.of(expression), Sources.NULL);
+    }
+
     public static UntypedStringLiteralNode stringLiteral(String value) {
         return new UntypedStringLiteralNode(value, Sources.NULL);
     }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionResumeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionResumeTests.java
new file mode 100644
index 00000000..c3326020
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionResumeTests.java
@@ -0,0 +1,116 @@
+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.CannotResumeOutsideOfEffectHandlerFieldError;
+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;
+
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.equalTo;
+
+public class TypeCheckerFunctionResumeTests {
+    @Test
+    public void expressionIsTypeChecked() {
+        var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.STRING));
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedResumeNode(isTypedStringLiteralNode("hello")));
+    }
+
+    @Test
+    public void whenExpressionIsAbsentThenUnitIsReturned() {
+        var untyped = Untyped.resume();
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.UNIT));
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedResumeNode(isTypedUnitLiteralNode()));
+    }
+
+    @Test
+    public void whenNotInFunctionThenResumeIsForbidden() {
+        var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespace(NamespaceName.of("Example"), P.map());
+
+        assertThrows(
+            CannotResumeOutsideOfEffectHandlerFieldError.class,
+            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+        );
+    }
+
+    @Test
+    public void whenInFunctionThenResumeIsForbidden() {
+        var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.explicit(Types.STRING));
+
+        assertThrows(
+            CannotResumeOutsideOfEffectHandlerFieldError.class,
+            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+        );
+    }
+
+    @Test
+    public void canResumeWIthSubtypeOfResumeType() {
+        var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.TOP));
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedResumeNode(isTypedStringLiteralNode("hello")));
+    }
+
+    @Test
+    public void cannotResumeWithNonSubtypeOfResumeType() {
+        var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.INT));
+
+        var error = assertThrows(
+            UnexpectedTypeError.class,
+            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+        );
+
+        assertThat(error.expected(), equalTo(TypeSets.type(Types.INT)));
+        assertThat(error.actual(), equalTo(Types.STRING));
+    }
+
+    @Test
+    public void resumeTypeIsUsedAsTypeHintForExpression() {
+        var untyped = Untyped.resume(
+            UntypedFunctionExpressionNode.builder()
+                .addPositionalParam("x", Optional.empty())
+                .withoutReturnType()
+                .build()
+        );
+        var resumeType = FunctionType.builder()
+            .addPositionalParam(Types.STRING)
+            .withReturnType(Types.UNIT)
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(resumeType));
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedResumeNode(
+            isTypedFunctionExpressionNode().withType(resumeType)
+        ));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
index f0c2daa7..8fe03d32 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
@@ -59,7 +59,7 @@ public class TypeCheckerFunctionReturnTests {
     public void whenInEffectHandlerThenReturnIsForbidden() {
         var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.UNIT));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.STRING));
 
         assertThrows(
             CannotReturnFromEffectHandlerFieldError.class,
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java
index cffae60a..169816dc 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java
@@ -366,7 +366,7 @@ public class TypeCheckerTryTests {
             .addHandlerField(
                 UntypedFunctionDefinitionNode.builder()
                     .withName("effectField")
-                    .addStatement(Untyped.expressionStatement(
+                    .addStatement(Untyped.resume(
                         UntypedCallNode.builder()
                             .withReceiver(Untyped.reference("f"))
                             .withHasEffects(true)