Clarify typing around handler methods

8 files changed, 190 insertions, 26 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 1c0cad87..faf6f31c 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -328,6 +328,10 @@ public class TypedNodeMatchers {
         return TypedTryNodeMatcher.empty();
     }
 
+    public static TypedTryHandlerMethodNodeMatcher isTypedTryHandlerMethodNode() {
+        return TypedTryHandlerMethodNodeMatcher.empty();
+    }
+
     public static Matcher<TypedExpressionNode> isTypedTupleNode(
         Matcher<Iterable<? extends TypedExpressionNode>> elements
     ) {
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNodeMatcher.java
new file mode 100644
index 00000000..bd19a7d6
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNodeMatcher.java
@@ -0,0 +1,58 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.pcollections.Empty;
+import org.pcollections.PVector;
+import org.shedlang.compiler.testing.MatcherBuilder;
+import org.shedlang.compiler.types.FunctionType;
+import org.zwobble.precisely.Matcher;
+
+import static org.zwobble.precisely.Matchers.equalTo;
+import static org.zwobble.precisely.Matchers.has;
+
+public class TypedTryHandlerMethodNodeMatcher extends MatcherBuilder<
+    TypedTryHandlerMethodNode,
+    TypedTryHandlerMethodNode,
+    TypedTryHandlerMethodNodeMatcher
+> {
+    public static TypedTryHandlerMethodNodeMatcher empty() {
+        return new TypedTryHandlerMethodNodeMatcher(Empty.vector());
+    }
+
+    public TypedTryHandlerMethodNodeMatcher(
+        PVector<Matcher<? super TypedTryHandlerMethodNode>> submatchers
+    ) {
+        super(submatchers);
+    }
+
+    @Override
+    protected Class<TypedTryHandlerMethodNode> expectedType() {
+        return TypedTryHandlerMethodNode.class;
+    }
+
+    @Override
+    protected TypedTryHandlerMethodNodeMatcher create(
+        PVector<Matcher<? super TypedTryHandlerMethodNode>> submatchers
+    ) {
+        return new TypedTryHandlerMethodNodeMatcher(submatchers);
+    }
+
+    public TypedTryHandlerMethodNodeMatcher withBodyBlock(Matcher<TypedBlockNode> body) {
+        return with(has("body", x -> x.body().block(), body));
+    }
+
+    public TypedTryHandlerMethodNodeMatcher withFreeVariables(Matcher<Iterable<FreeVariable>> freeVariables) {
+        return with(has("freeVariables", x -> x.body().freeVariables(), freeVariables));
+    }
+
+    public TypedTryHandlerMethodNodeMatcher withName(String name) {
+        return with(has("name", x -> x.name(), equalTo(name)));
+    }
+
+    public TypedTryHandlerMethodNodeMatcher withParams(Matcher<TypedParamsNode> params) {
+        return with(has("params", x -> x.params(), params));
+    }
+
+    public TypedTryHandlerMethodNodeMatcher withType(FunctionType type) {
+        return with(has("type", x -> x.type(), equalTo(type)));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedTryNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedTryNodeMatcher.java
index 6998f85a..9bc549d3 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedTryNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedTryNodeMatcher.java
@@ -39,7 +39,7 @@ public class TypedTryNodeMatcher extends MatcherBuilder<
     }
 
     public TypedTryNodeMatcher withHandlerMethods(
-        Matcher<? super Iterable<TypedFunctionDefinitionNode>> handlerMethods
+        Matcher<? super Iterable<TypedTryHandlerMethodNode>> handlerMethods
     ) {
         return with(has("handlerMethods", x -> x.handlerMethods(), handlerMethods));
     }
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
index d2df2a3c..41c51713 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
@@ -324,6 +324,10 @@ public class UntypedNodeMatchers {
         return UntypedTryNodeMatcher.empty();
     }
 
+    public static UntypedTryHandlerMethodNodeMatcher isUntypedTryHandlerMethodNode() {
+        return UntypedTryHandlerMethodNodeMatcher.empty();
+    }
+
     public static Matcher<UntypedExpressionNode> isUntypedTupleNode(
         Matcher<Iterable<? extends UntypedExpressionNode>> elements
     ) {
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryHandlerMethodNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryHandlerMethodNodeMatcher.java
new file mode 100644
index 00000000..8e08587d
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryHandlerMethodNodeMatcher.java
@@ -0,0 +1,51 @@
+package org.shedlang.compiler.ast.untyped;
+
+import org.pcollections.Empty;
+import org.pcollections.PVector;
+import org.shedlang.compiler.testing.MatcherBuilder;
+import org.zwobble.precisely.Matcher;
+
+import static org.zwobble.precisely.Matchers.equalTo;
+import static org.zwobble.precisely.Matchers.has;
+
+public class UntypedTryHandlerMethodNodeMatcher extends MatcherBuilder<
+    UntypedTryHandlerMethodNode,
+    UntypedTryHandlerMethodNode,
+    UntypedTryHandlerMethodNodeMatcher
+> {
+    public static UntypedTryHandlerMethodNodeMatcher empty() {
+        return new UntypedTryHandlerMethodNodeMatcher(Empty.vector());
+    }
+
+    public UntypedTryHandlerMethodNodeMatcher(
+        PVector<Matcher<? super UntypedTryHandlerMethodNode>> submatchers
+    ) {
+        super(submatchers);
+    }
+
+    @Override
+    protected Class<UntypedTryHandlerMethodNode> expectedType() {
+        return UntypedTryHandlerMethodNode.class;
+    }
+
+    @Override
+    protected UntypedTryHandlerMethodNodeMatcher create(
+        PVector<Matcher<? super UntypedTryHandlerMethodNode>> submatchers
+    ) {
+        return new UntypedTryHandlerMethodNodeMatcher(submatchers);
+    }
+
+    public UntypedTryHandlerMethodNodeMatcher withParams(
+        Matcher<? super UntypedParamsNode> params
+    ) {
+        return with(has("params", x -> x.params(), params));
+    }
+
+    public UntypedTryHandlerMethodNodeMatcher withBody(Matcher<UntypedBlockNode> body) {
+        return with(has("body", x -> x.body(), body));
+    }
+
+    public UntypedTryHandlerMethodNodeMatcher withName(String name) {
+        return with(has("name", x -> x.name(), equalTo(name)));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryNodeMatcher.java
index eac572ab..b5df5392 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedTryNodeMatcher.java
@@ -43,7 +43,7 @@ public class UntypedTryNodeMatcher extends MatcherBuilder<
     }
 
     public UntypedTryNodeMatcher withHandlerMethods(
-        Matcher<Iterable<? extends UntypedFunctionDefinitionNode>> handlerMethods
+        Matcher<Iterable<? extends UntypedTryHandlerMethodNode>> handlerMethods
     ) {
         return with(has("handlerMethods", x -> x.handlerMethods(), handlerMethods));
     }
diff --git a/src/test/java/org/shedlang/compiler/parser/ParserTryTests.java b/src/test/java/org/shedlang/compiler/parser/ParserTryTests.java
index 0782cdd4..3a1a2ac2 100644
--- a/src/test/java/org/shedlang/compiler/parser/ParserTryTests.java
+++ b/src/test/java/org/shedlang/compiler/parser/ParserTryTests.java
@@ -5,7 +5,7 @@ import org.junit.jupiter.api.Test;
 import static org.shedlang.compiler.ast.untyped.UntypedNodeMatchers.*;
 import static org.shedlang.compiler.parser.ParserTesting.parseString;
 import static org.zwobble.precisely.AssertThat.assertThat;
-import static org.zwobble.precisely.Matchers.isSequence;
+import static org.zwobble.precisely.Matchers.*;
 
 public class ParserTryTests {
     @Test
@@ -15,9 +15,10 @@ public class ParserTryTests {
                 1;
             } handle State {
                 fun read() {
+                    resume 42;
                 }
                 
-                fun write() {
+                fun write(value) {
                 }
             }
         """;
@@ -32,10 +33,21 @@ public class ParserTryTests {
                 )))
                 .withHandledEffect(isUntypedTypeLevelReferenceNode("State"))
                 .withHandlerMethods(isSequence(
-                    isUntypedFunctionDefinitionNode()
-                        .withName("read"),
-                    isUntypedFunctionDefinitionNode()
+                    isUntypedTryHandlerMethodNode()
+                        .withName("read")
+                        .withParams(isUntypedParamsNode(isSequence(), isSequence()))
+                        .withBody(isUntypedBlockNode(isSequence(
+                            isUntypedResumeNode(isOptionalOf(isUntypedIntLiteralNode(42)))
+                        ))),
+                    isUntypedTryHandlerMethodNode()
                         .withName("write")
+                        .withParams(isUntypedParamsNode(
+                            isSequence(
+                                isUntypedParamNode("value", isOptionalEmpty())
+                            ),
+                            isSequence()
+                        ))
+                        .withBody(isUntypedBlockNode(isSequence()))
                 ))
         );
     }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java
index f62780d4..c23511be 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerTryTests.java
@@ -1,10 +1,7 @@
 package org.shedlang.compiler.typechecker;
 
 import org.junit.jupiter.api.Test;
-import org.shedlang.compiler.ast.untyped.Untyped;
-import org.shedlang.compiler.ast.untyped.UntypedCallNode;
-import org.shedlang.compiler.ast.untyped.UntypedFunctionDefinitionNode;
-import org.shedlang.compiler.ast.untyped.UntypedTryNode;
+import org.shedlang.compiler.ast.untyped.*;
 import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.errors.NoSuchFieldError;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
@@ -197,7 +194,7 @@ public class TypeCheckerTryTests {
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .withName("mark")
                     .build()
             )
@@ -230,7 +227,7 @@ public class TypeCheckerTryTests {
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .addPositionalParam("x", Untyped.typeLevelReference("String"))
                     .addStatement(Untyped.resume())
                     .withName("mark")
@@ -277,7 +274,7 @@ public class TypeCheckerTryTests {
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .addStatement(Untyped.resume())
                     .withName("mark")
                     .build()
@@ -304,7 +301,7 @@ public class TypeCheckerTryTests {
             result.value(),
             isTypedTryNode()
                 .withHandlerMethods(isSequence(
-                    isTypedFunctionDefinitionNode()
+                    isTypedTryHandlerMethodNode()
                         .withName("mark")
                         .withBodyBlock(isTypedBlockNode(isSequence(
                             isTypedResumeNode(isTypedUnitLiteralNode())
@@ -314,14 +311,14 @@ public class TypeCheckerTryTests {
     }
 
     @Test
-    public void whenHandlerMethodTypeHasWrongReturnTypeThenErrorIsThrown() {
+    public void whenHandlerMethodTypeHasWrongResumeTypeThenErrorIsThrown() {
         var handledEffect = Effects.userDefined(NAMESPACE_NAME, "HandledEffect");
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .addStatement(Untyped.resume())
-                    .withReturnType(Untyped.typeLevelReference("String"))
+                    .withResumeType(Untyped.typeLevelReference("String"))
                     .withName("mark")
                     .build()
             )
@@ -352,14 +349,14 @@ public class TypeCheckerTryTests {
     }
 
     @Test
-    public void whenHandlerMethodTypeHasNarrowerReturnTypeThanEffectMethodThenNarrowerTypeIsUsedToTypeCheckBody() {
+    public void whenHandlerMethodTypeHasNarrowerResumeTypeThanEffectMethodThenNarrowerTypeIsUsedToTypeCheckBody() {
         var handledEffect = Effects.userDefined(NAMESPACE_NAME, "HandledEffect");
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .addStatement(Untyped.resume(Untyped.intLiteral(0)))
-                    .withReturnType(Untyped.typeLevelReference("String"))
+                    .withResumeType(Untyped.typeLevelReference("String"))
                     .withName("mark")
                     .build()
             )
@@ -390,12 +387,12 @@ public class TypeCheckerTryTests {
     }
 
     @Test
-    public void whenHandlerMethodTypeHasEffectThenEffectOfExpressionScopeIsIgnored() {
+    public void whenHandlerMethodTypeHasNarrowerEffectThanExpressionScopeThenEffectOfExpressionScopeIsIgnored() {
         var handledEffect = Effects.userDefined(NAMESPACE_NAME, "HandledEffect");
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .withName("effectField")
                     .withEffect(Untyped.typeLevelReference("Pure"))
                     .addStatement(Untyped.resume(
@@ -434,13 +431,51 @@ public class TypeCheckerTryTests {
     }
 
     @Test
+    public void whenHandlerMethodTypeHasNonSubEffectOfExpressionScopeThenErrorIsThrown() {
+        var handledEffect = Effects.userDefined(NAMESPACE_NAME, "HandledEffect");
+        var untyped = UntypedTryNode.builder()
+            .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
+            .addHandlerMethod(
+                UntypedTryHandlerMethodNode.builder()
+                    .withName("effectField")
+                    .withEffect(Untyped.typeLevelReference("Io"))
+                    .addStatement(Untyped.resume())
+                    .build()
+            )
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(
+                P.map(
+                    Map.entry(handledEffect.name(), Variable.defined(Types.effectType(handledEffect), VariableType.CONSTANT)),
+                    Map.entry("Io", Variable.defined(Types.effectType(Effects.IO), VariableType.CONSTANT))
+                ),
+                Effects.NONE,
+                Optional.empty()
+            )
+            .addUserDefinedEffectInfo(
+                handledEffect,
+                UserDefinedEffectInfo.builder()
+                    .addField("effectField", FunctionType.builder().withReturnType(Types.UNIT).build())
+                    .build()
+            );
+
+        var error = assertThrows(
+            UnexpectedEffectError.class,
+            () -> typeCheckExpression(untyped, context)
+        );
+
+        assertThat(error.allowed(), equalTo(Effects.NONE));
+        assertThat(error.actual(), equalTo(Effects.IO));
+    }
+
+    @Test
     public void whenEffectIsNotAllowedInExpressionScopeAndIsNotHandledEffectThenEffectIsNotAllowedInHandlerScope() {
         var handledEffect = Effects.userDefined(NAMESPACE_NAME, "HandledEffect");
         var unexpectedEffect = Effects.userDefined(NAMESPACE_NAME, "UnexpectedEffect");
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .withName("effectField")
                     .addStatement(Untyped.expressionStatement(
                         UntypedCallNode.builder()
@@ -484,7 +519,7 @@ public class TypeCheckerTryTests {
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .withName("effectField")
                     .addStatement(Untyped.resume(
                         UntypedCallNode.builder()
@@ -523,7 +558,7 @@ public class TypeCheckerTryTests {
         var untyped = UntypedTryNode.builder()
             .withHandledEffect(Untyped.typeLevelReference(handledEffect.name()))
             .addHandlerMethod(
-                UntypedFunctionDefinitionNode.builder()
+                UntypedTryHandlerMethodNode.builder()
                     .withName("effectField")
                     .addStatement(Untyped.expressionStatement(
                         UntypedCallNode.builder()