Store handled effect in typed break node

3 files changed, 18 insertions, 5 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(