Store handled effect in typed break node

7 files changed, 50 insertions, 23 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java
index bc7ea34a..16e29997 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java
@@ -1,8 +1,10 @@
 package org.shedlang.compiler.ast.typed;
 
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.UserDefinedEffect;
 
 public record TypedBreakNode(
+    UserDefinedEffect handledEffect,
     TypedExpressionNode expression,
     Source source
 ) implements TypedFunctionStatementNode {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
index 201489d5..be31d0a7 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.typechecker;
 
 import org.pcollections.PVector;
 import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.types.UserDefinedEffect;
 import org.shedlang.compiler.util.P;
 
 public sealed interface ReturnTypeState {
@@ -16,8 +17,11 @@ public sealed interface ReturnTypeState {
         return new ImplicitLowerBound(type);
     }
 
-    static ReturnTypeState effectHandler(Type resumeType) {
-        return new EffectHandler(resumeType, P.vector());
+    static ReturnTypeState effectHandler(
+        UserDefinedEffect handledEffect,
+        Type resumeType
+    ) {
+        return new EffectHandler(handledEffect, resumeType, P.vector());
     }
 
     record NotInFunction() implements ReturnTypeState {}
@@ -25,9 +29,13 @@ public sealed interface ReturnTypeState {
     record ImplicitUnknown() implements ReturnTypeState {}
     record ImplicitLowerBound(Type type) implements ReturnTypeState {}
 
-    record EffectHandler(Type resumeType, PVector<Type> breakTypes) implements ReturnTypeState {
+    record EffectHandler(
+        UserDefinedEffect handledEffect,
+        Type resumeType,
+        PVector<Type> breakTypes
+    ) implements ReturnTypeState {
         public ReturnTypeState addBreakType(Type breakType) {
-            return new EffectHandler(resumeType, breakTypes.plus(breakType));
+            return new EffectHandler(handledEffect, resumeType, breakTypes.plus(breakType));
         }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 94ba7ea8..65046b42 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -309,6 +309,7 @@ public class TypeChecker {
                 );
 
                 return context.toResult(new TypedBreakNode(
+                    effectHandler.handledEffect(),
                     typedExpression,
                     untyped.source()
                 ));
@@ -1754,6 +1755,7 @@ public class TypeChecker {
 
             var typedHandlerMethodResult = typeCheckTryHandlerMethod(
                 untypedHandlerMethod,
+                handledEffect,
                 effectFieldType,
                 context
             );
@@ -1794,6 +1796,7 @@ public class TypeChecker {
 
     private static TypeCheckerResult<TypedTryHandlerMethodNode> typeCheckTryHandlerMethod(
         UntypedTryHandlerMethodNode untyped,
+        UserDefinedEffect handledEffect,
         FunctionType effectFieldType,
         TypeCheckerContext context
     ) {
@@ -1837,7 +1840,7 @@ public class TypeChecker {
             inferredParamTypes.positionalParams,
             inferredParamTypes.namedParams,
             effect,
-            ReturnTypeState.effectHandler(resumeType)
+            ReturnTypeState.effectHandler(handledEffect, resumeType)
         );
 
         var typedBodyResult = typeCheckFunctionBody(
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index 0014e4f0..4a5bf1c5 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -69,10 +69,12 @@ public class TypedNodeMatchers {
     }
 
     public static Matcher<TypedFunctionStatementNode> isTypedBreakNode(
+        UserDefinedEffect handledEffect,
         Matcher<TypedExpressionNode> expression
     ) {
         return instanceOf(
             TypedBreakNode.class,
+            has("handledEffect", x -> x.handledEffect(), equalTo(handledEffect)),
             has("expression", x -> x.expression(), expression)
         );
     }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
index 32f7b127..c355b0cb 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
@@ -6,6 +6,7 @@ import org.shedlang.compiler.typechecker.errors.CannotBreakOutsideOfEffectHandle
 import org.shedlang.compiler.types.Effects;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Types;
+import org.shedlang.compiler.types.UserDefinedEffect;
 import org.shedlang.compiler.util.P;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -13,26 +14,35 @@ import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
 import static org.zwobble.precisely.AssertThat.assertThat;
 
 public class TypeCheckerBreakTests {
+    private static final NamespaceName NAMESPACE_NAME = NamespaceName.of("Example");
+    private static final UserDefinedEffect EFFECT = Effects.userDefined(NAMESPACE_NAME, "Effect");
+
     @Test
     public void expressionIsTypeChecked() {
         var untyped = Untyped.break_(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.INT));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.INT));
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
-        assertThat(typed.value(), isTypedBreakNode(isTypedStringLiteralNode("hello")));
+        assertThat(typed.value(), isTypedBreakNode(
+            EFFECT,
+            isTypedStringLiteralNode("hello")
+        ));
     }
 
     @Test
     public void whenExpressionIsAbsentThenUnitIsReturned() {
         var untyped = Untyped.break_();
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.UNIT));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.UNIT));
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
-        assertThat(typed.value(), isTypedBreakNode(isTypedUnitLiteralNode()));
+        assertThat(typed.value(), isTypedBreakNode(
+            EFFECT,
+            isTypedUnitLiteralNode()
+        ));
     }
 
     @Test
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
index f74b6d26..6d8c75df 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionReturnTests.java
@@ -7,10 +7,7 @@ import org.shedlang.compiler.typechecker.errors.CannotReturnFromEffectHandlerMet
 import org.shedlang.compiler.typechecker.errors.CannotReturnOutsideOfFunctionError;
 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.types.*;
 import org.shedlang.compiler.util.P;
 
 import java.util.Optional;
@@ -58,8 +55,13 @@ public class TypeCheckerFunctionReturnTests {
     @Test
     public void whenInEffectHandlerThenReturnIsForbidden() {
         var untyped = Untyped.functionReturn(Untyped.stringLiteral("hello"));
+        var effect = Effects.userDefined(NamespaceName.of("Example"), "State");
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.STRING));
+            .enterFunction(
+                P.map(),
+                Effects.NONE,
+                ReturnTypeState.effectHandler(effect, Types.STRING)
+            );
 
         assertThrows(
             CannotReturnFromEffectHandlerMethodError.class,
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
index b5bab880..91ba8c74 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
@@ -6,10 +6,7 @@ import org.shedlang.compiler.ast.untyped.UntypedFunctionExpressionNode;
 import org.shedlang.compiler.typechecker.errors.CannotResumeOutsideOfEffectHandlerMethodError;
 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.types.*;
 import org.shedlang.compiler.util.P;
 
 import java.util.Optional;
@@ -20,11 +17,14 @@ import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.equalTo;
 
 public class TypeCheckerResumeTests {
+    private static final NamespaceName NAMESPACE_NAME = NamespaceName.of("Example");
+    private static final UserDefinedEffect EFFECT = Effects.userDefined(NAMESPACE_NAME, "Effect");
+
     @Test
     public void expressionIsTypeChecked() {
         var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.STRING));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.STRING));
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
@@ -35,7 +35,7 @@ public class TypeCheckerResumeTests {
     public void whenExpressionIsAbsentThenUnitIsReturned() {
         var untyped = Untyped.resume();
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.UNIT));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.UNIT));
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
@@ -70,7 +70,7 @@ public class TypeCheckerResumeTests {
     public void canResumeWithSubtypeOfResumeType() {
         var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.TOP));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.TOP));
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
 
@@ -81,7 +81,7 @@ public class TypeCheckerResumeTests {
     public void cannotResumeWithNonSubtypeOfResumeType() {
         var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.INT));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, Types.INT));
 
         var error = assertThrows(
             UnexpectedTypeError.class,
@@ -105,7 +105,7 @@ public class TypeCheckerResumeTests {
             .withReturnType(Types.UNIT)
             .build();
         var context = TypeCheckerContext.EMPTY
-            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(resumeType));
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(EFFECT, resumeType));
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);