Allow capture of vars in effect handler methods in type checker

2 files changed, 18 insertions, 12 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
index cfcd59ec..e8e86885 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
@@ -23,6 +23,10 @@ public record PartialFunctionSignature(
         return Lists.concat(positionalParams, namedParams);
     }
 
+    public boolean isEffectHandlerMethod() {
+        return this.returnTypeState instanceof ReturnTypeState.EffectHandler;
+    }
+
     public PartialFunctionSignature withEffect(Effect effect) {
         return new PartialFunctionSignature(
             this.positionalParams,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 0bb29d5f..3047285d 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -644,18 +644,20 @@ public class TypeChecker {
                     freeVariable.firstReferenceSource()
                 ));
 
-            switch (variable.variableType()) {
-                case CONSTANT -> {}
-                case IMPORT -> {}
-                case PARAM -> {}
-                case VAL -> {}
-
-                case VAR ->
-                    throw new CannotCaptureVarError(
-                        freeVariable.name(),
-                        freeVariable.firstReferenceSource()
-                    );
-            };
+            if (!signature.isEffectHandlerMethod()) {
+                switch (variable.variableType()) {
+                    case CONSTANT -> {}
+                    case IMPORT -> {}
+                    case PARAM -> {}
+                    case VAL -> {}
+
+                    case VAR ->
+                        throw new CannotCaptureVarError(
+                            freeVariable.name(),
+                            freeVariable.firstReferenceSource()
+                        );
+                };
+            }
         }
 
         var typed = new TypedFunctionBodyNode(