Store handled effect in typed break node

4 files changed, 32 insertions, 18 deletions

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);