Type check resume statements properly

3 files changed, 59 insertions, 7 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);
+    }
+}