Use explicit resume type

1 files changed, 24 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 3bc3e088..c45b770b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1737,9 +1737,30 @@ public class TypeChecker {
                 contextForReturnType
             ));
 
-        var partialSignature = inferFunctionPartialSignature(typedSignature, Optional.of(effectFieldType))
-            // TODO: use explicit return type if present
-            .withReturnTypeState(ReturnTypeState.effectHandler(effectFieldType.returnType()));
+        var partialSignature = inferFunctionPartialSignature(typedSignature, Optional.of(effectFieldType));
+
+        // TODO: Find a nicer way of doing this (we know the return type state should only be one of two variants).
+        var effectHandlerReturnTypeState = switch (partialSignature.returnTypeState()) {
+            case ReturnTypeState.EffectHandler effectHandler ->
+                throw new IllegalStateException();
+
+            case ReturnTypeState.Explicit explicit ->
+                // TODO: do we actually want to allow this? Potentially unclear
+                // whether this refers to the resume type or the exit type?
+                ReturnTypeState.effectHandler(explicit.type());
+
+            case ReturnTypeState.ImplicitLowerBound implicitLowerBound ->
+                throw new IllegalStateException();
+
+            case ReturnTypeState.ImplicitUnknown implicitUnknown ->
+                ReturnTypeState.effectHandler(effectFieldType.returnType());
+
+            case ReturnTypeState.NotInFunction notInFunction ->
+                throw new IllegalStateException();
+
+        };
+
+        partialSignature = partialSignature.withReturnTypeState(effectHandlerReturnTypeState);
 
         if (typedSignature.effect().isEmpty()) {
             partialSignature = partialSignature.withEffect(context.effect());