Clarify typing around handler methods

6 files changed, 316 insertions, 55 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java
new file mode 100644
index 00000000..14253c67
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java
@@ -0,0 +1,118 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.pcollections.Empty;
+import org.pcollections.PVector;
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
+import org.shedlang.compiler.types.*;
+
+import java.util.List;
+
+public record TypedTryHandlerMethodNode(
+    String name,
+    TypedParamsNode params,
+    TypedFunctionBodyNode body,
+    FunctionType type,
+    Source source
+) implements TypedNode {
+    public static Builder builder() {
+        return new Builder(
+            "f",
+            TypedFunctionSignatureNode.builder(),
+            FreeVariables.builder(),
+            Empty.vector(),
+            Empty.vector(),
+            new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT)
+        );
+    }
+
+    public record Builder(
+        String name,
+        TypedFunctionSignatureNode.Builder signature,
+        FreeVariables.Builder freeVariables,
+        PVector<String> boundVariableNames,
+        PVector<TypedFunctionStatementNode> body,
+        FunctionType type
+    ) {
+        public TypedTryHandlerMethodNode build() {
+            return new TypedTryHandlerMethodNode(
+                name,
+                signature.build().params(),
+                new TypedFunctionBodyNode(
+                    freeVariables.build(),
+                    new TypedBlockNode(body, boundVariableNames, Sources.NULL),
+                    Types.UNIT,
+                    Sources.NULL
+                ),
+                type,
+                Sources.NULL
+            );
+        }
+
+        public TypedTryHandlerMethodNode.Builder withName(String name) {
+            return new TypedTryHandlerMethodNode.Builder(
+                name,
+                this.signature,
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body,
+                this.type
+            );
+        }
+
+        public TypedTryHandlerMethodNode.Builder addPositionalParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new TypedTryHandlerMethodNode.Builder(
+                this.name,
+                this.signature.addPositionalParam(name, type),
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body,
+                this.type
+            ).addBoundVariable(name);
+        }
+
+        public TypedTryHandlerMethodNode.Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new TypedTryHandlerMethodNode.Builder(
+                this.name,
+                this.signature.addNamedParam(name, type),
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body,
+                this.type
+            ).addBoundVariable(name);
+        }
+
+        public Builder addFreeVariable(String freeVariable) {
+            return new TypedTryHandlerMethodNode.Builder(
+                this.name,
+                this.signature,
+                this.freeVariables.plus(freeVariable, Sources.NULL),
+                this.boundVariableNames,
+                this.body,
+                this.type
+            );
+        }
+
+        public Builder addBoundVariable(String name) {
+            return new TypedTryHandlerMethodNode.Builder(
+                this.name,
+                this.signature,
+                this.freeVariables,
+                this.boundVariableNames.plus(name),
+                this.body,
+                this.type
+            );
+        }
+
+        public TypedTryHandlerMethodNode.Builder addStatement(TypedFunctionStatementNode statement) {
+            return new TypedTryHandlerMethodNode.Builder(
+                this.name,
+                this.signature,
+                this.freeVariables,
+                this.boundVariableNames,
+                this.body.plus(statement),
+                this.type
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
index b58b0b3a..123744d5 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
@@ -9,7 +9,7 @@ import java.util.List;
 public record TypedTryNode(
     TypedBlockNode body,
     TypedTypeLevelExpressionNode<Effect> handledEffect,
-    List<TypedFunctionDefinitionNode> handlerMethods,
+    List<TypedTryHandlerMethodNode> handlerMethods,
     Source source
 ) implements TypedExpressionNode {
     @Override
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryHandlerMethodNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryHandlerMethodNode.java
new file mode 100644
index 00000000..193fdf49
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryHandlerMethodNode.java
@@ -0,0 +1,98 @@
+package org.shedlang.compiler.ast.untyped;
+
+import org.pcollections.Empty;
+import org.pcollections.PVector;
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
+
+import java.util.Optional;
+
+public record UntypedTryHandlerMethodNode(
+    String name,
+    UntypedParamsNode params,
+    Optional<UntypedTypeLevelExpressionNode> effect,
+    Optional<UntypedTypeLevelExpressionNode> resumeType,
+    UntypedBlockNode body,
+    Source source
+) implements UntypedNode {
+    public static Builder builder() {
+        return new Builder(
+            "f",
+            // TODO: separate params builder
+            UntypedFunctionSignatureNode.builder(),
+            Empty.vector()
+        );
+    }
+
+    public record Builder(
+        String name,
+        UntypedFunctionSignatureNode.Builder signature,
+        PVector<UntypedFunctionStatementNode> body
+    ) {
+        public UntypedTryHandlerMethodNode build() {
+            return new UntypedTryHandlerMethodNode(
+                name,
+                signature.build().params(),
+                signature.build().effect(),
+                signature.build().returnType(),
+                new UntypedBlockNode(body, Sources.NULL),
+                Sources.NULL
+            );
+        }
+
+        public Builder withName(String name) {
+            return new Builder(
+                name,
+                this.signature,
+                this.body
+            );
+        }
+
+        public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
+            return addPositionalParam(name, Optional.of(type));
+        }
+
+        public Builder addPositionalParam(
+            String name,
+            Optional<UntypedTypeLevelExpressionNode> type
+        ) {
+            return new Builder(
+                this.name,
+                this.signature.addPositionalParam(name, type),
+                this.body
+            );
+        }
+
+        public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
+            return new Builder(
+                this.name,
+                this.signature.addNamedParam(name, type),
+                this.body
+            );
+        }
+
+        public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
+            return new Builder(
+                this.name,
+                this.signature.withEffect(effect),
+                this.body
+            );
+        }
+
+        public Builder withResumeType(UntypedTypeLevelExpressionNode resumeType) {
+            return new Builder(
+                this.name,
+                this.signature.withReturnType(resumeType),
+                this.body
+            );
+        }
+
+        public Builder addStatement(UntypedFunctionStatementNode statement) {
+            return new Builder(
+                this.name,
+                this.signature.addStatement(statement),
+                this.body.plus(statement)
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryNode.java
index a411039e..95e85a61 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTryNode.java
@@ -10,7 +10,7 @@ import java.util.List;
 public record UntypedTryNode(
     UntypedBlockNode body,
     UntypedTypeLevelExpressionNode handledEffect,
-    List<UntypedFunctionDefinitionNode> handlerMethods,
+    List<UntypedTryHandlerMethodNode> handlerMethods,
     Source source
 ) implements UntypedExpressionNode {
     public static Builder builder() {
@@ -24,7 +24,7 @@ public record UntypedTryNode(
     public record Builder(
         PVector<UntypedFunctionStatementNode> body,
         UntypedTypeLevelExpressionNode handledEffect,
-        PVector<UntypedFunctionDefinitionNode> handlerMethods
+        PVector<UntypedTryHandlerMethodNode> handlerMethods
     ) {
         public UntypedTryNode build() {
             return new UntypedTryNode(
@@ -51,7 +51,7 @@ public record UntypedTryNode(
             );
         }
 
-        public Builder addHandlerMethod(UntypedFunctionDefinitionNode handlerMethod) {
+        public Builder addHandlerMethod(UntypedTryHandlerMethodNode handlerMethod) {
             return new Builder(
                 this.body,
                 this.handledEffect,
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index 55e2202d..72bbbb32 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -1097,7 +1097,7 @@ public class Parser {
 
         var handlerMethods = parseMany(
             () -> tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE),
-            () -> parseFunctionDefinition(tokens),
+            () -> parseTryHandlerMethod(tokens),
             () -> true
         );
 
@@ -1109,6 +1109,31 @@ public class Parser {
         return new UntypedTryNode(body, handledEffect, handlerMethods, source);
     }
 
+    private static UntypedTryHandlerMethodNode parseTryHandlerMethod(TokenIterator<TokenType> tokens) {
+        var start = tokens.position();
+
+        tokens.skip(TokenType.KEYWORD_FUN);
+        var name = tokens.nextValue(TokenType.IDENTIFIER);
+
+        tokens.skip(TokenType.SYMBOL_PAREN_OPEN);
+        var params = parseParams(tokens);
+        tokens.skip(TokenType.SYMBOL_PAREN_CLOSE);
+
+        var body = parseBlock(tokens);
+
+        var end = tokens.position();
+        var source = start.to(end);
+
+        return new UntypedTryHandlerMethodNode(
+            name,
+            params,
+            Optional.empty(),
+            Optional.empty(),
+            body,
+            source
+        );
+    }
+
     private static UntypedExpressionNode parseTuple(
         TokenIterator<TokenType> tokens
     ) {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index c45b770b..5c203322 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -884,17 +884,25 @@ public class TypeChecker {
         );
     }
 
-    private static PartialFunctionSignature inferFunctionPartialSignature(
-        TypedFunctionSignatureNode typedSignature,
+    // TODO: lift out general Params record?
+    private record InferredParamsTypes(
+        List<PartialFunctionSignature.Param> positionalParams,
+        List<PartialFunctionSignature.Param> namedParams
+    ) {
+
+    }
+
+    private static InferredParamsTypes inferParamTypes(
+        TypedParamsNode params,
         Optional<Type> hint
     ) {
         var functionTypeHint = Optionals.filterToType(FunctionType.class, hint);
 
         var positionalParams = new ArrayList<PartialFunctionSignature.Param>();
 
-        for (var i = 0; i < typedSignature.params().positional().size(); i++) {
+        for (var i = 0; i < params.positional().size(); i++) {
             var paramIndex = i;
-            var paramNode = typedSignature.params().positional().get(paramIndex);
+            var paramNode = params.positional().get(paramIndex);
             var paramHint = functionTypeHint.flatMap(h -> Lists.tryGet(h.positionalParams(), paramIndex));
 
             var param = new PartialFunctionSignature.Param(
@@ -910,7 +918,7 @@ public class TypeChecker {
 
         var namedParams = new ArrayList<PartialFunctionSignature.Param>();
 
-        for (var namedParam : typedSignature.params().named()) {
+        for (var namedParam : params.named()) {
             var paramHint = functionTypeHint.flatMap(
                 h -> h.namedParams().stream()
                     .filter(p -> p.name().equals(namedParam.name()))
@@ -928,13 +936,22 @@ public class TypeChecker {
             ));
         }
 
+        return new InferredParamsTypes(positionalParams, namedParams);
+    }
+
+    private static PartialFunctionSignature inferFunctionPartialSignature(
+        TypedFunctionSignatureNode typedSignature,
+        Optional<Type> hint
+    ) {
+        var inferredParamTypes = inferParamTypes(typedSignature.params(), hint);
+
         // TODO: if effect is not set, infer it from body, with an upper bound
         // of the context's effect.
         var effect = functionSignatureToEffect(typedSignature);
 
         return new PartialFunctionSignature(
-            positionalParams,
-            namedParams,
+            inferredParamTypes.positionalParams,
+            inferredParamTypes.namedParams,
             effect,
             typedSignature.returnType().map(node -> node.value())
                 .map(ReturnTypeState::explicit)
@@ -1680,7 +1697,7 @@ public class TypeChecker {
         var effectInfo = context.userDefinedEffectInfo(handledEffect, typedHandledEffect.source());
         var handledFieldNames = new HashSet<String>();
 
-        var typedHandlerMethods = new ArrayList<TypedFunctionDefinitionNode>();
+        var typedHandlerMethods = new ArrayList<TypedTryHandlerMethodNode>();
         for (var untypedHandlerMethod : untyped.handlerMethods()) {
             handledFieldNames.add(untypedHandlerMethod.name());
 
@@ -1721,50 +1738,53 @@ public class TypeChecker {
         ));
     }
 
-    private static TypeCheckerResult<TypedFunctionDefinitionNode> typeCheckTryHandlerMethod(
-        UntypedFunctionDefinitionNode untyped,
+    private static TypeCheckerResult<TypedTryHandlerMethodNode> typeCheckTryHandlerMethod(
+        UntypedTryHandlerMethodNode untyped,
         FunctionType effectFieldType,
         TypeCheckerContext context
     ) {
-        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
-
-        var contextForReturnType = context;
-        typedSignature.returnType()
-            .ifPresent(returnType -> checkIsSubType(
-                returnType.value(),
-                effectFieldType.returnType(),
-                returnType.source(),
-                contextForReturnType
-            ));
-
-        var partialSignature = inferFunctionPartialSignature(typedSignature, Optional.of(effectFieldType));
-
-        // TODO: Find a nicer way of doing this (we know the return type state should only be one of two variants).
-        var effectHandlerReturnTypeState = switch (partialSignature.returnTypeState()) {
-            case ReturnTypeState.EffectHandler effectHandler ->
-                throw new IllegalStateException();
-
-            case ReturnTypeState.Explicit explicit ->
-                // TODO: do we actually want to allow this? Potentially unclear
-                // whether this refers to the resume type or the exit type?
-                ReturnTypeState.effectHandler(explicit.type());
-
-            case ReturnTypeState.ImplicitLowerBound implicitLowerBound ->
-                throw new IllegalStateException();
-
-            case ReturnTypeState.ImplicitUnknown implicitUnknown ->
-                ReturnTypeState.effectHandler(effectFieldType.returnType());
-
-            case ReturnTypeState.NotInFunction notInFunction ->
-                throw new IllegalStateException();
-
-        };
+        var typedParams = typeCheckParams(untyped.params(), context);
+        var inferredParamTypes = inferParamTypes(typedParams, Optional.of(effectFieldType));
 
-        partialSignature = partialSignature.withReturnTypeState(effectHandlerReturnTypeState);
+        var initialContext = context;
+        var expectedEffect = initialContext.effect();
+        var effect = untyped.effect()
+            .map(untypedEffectNode -> {
+                var explicitEffect = typeCheckEffect(untypedEffectNode, initialContext).value();
+                if (!isSubEffect(explicitEffect, expectedEffect, initialContext)) {
+                    throw new UnexpectedEffectError(
+                        expectedEffect,
+                        explicitEffect,
+                        untypedEffectNode.source()
+                    );
+                }
+                return explicitEffect;
+            })
+            .orElse(expectedEffect);
+
+        var expectedResumeType = effectFieldType.returnType();
+        var resumeType = untyped.resumeType()
+            .map(untypedResumeTypeNode -> {
+                var explicitResumeType = typeCheckTypeExpression(
+                    untypedResumeTypeNode,
+                    initialContext
+                ).value();
+                checkIsSubType(
+                    explicitResumeType,
+                    expectedResumeType,
+                    untypedResumeTypeNode.source(),
+                    initialContext
+                );
+                return explicitResumeType;
+            })
+            .orElse(expectedResumeType);
 
-        if (typedSignature.effect().isEmpty()) {
-            partialSignature = partialSignature.withEffect(context.effect());
-        }
+        var partialSignature = new PartialFunctionSignature(
+            inferredParamTypes.positionalParams,
+            inferredParamTypes.namedParams,
+            effect,
+            ReturnTypeState.effectHandler(resumeType)
+        );
 
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
@@ -1782,12 +1802,12 @@ public class TypeChecker {
                 .map(x -> new NamedParam(x.name(), x.type()))
                 .toList(),
             partialSignature.effect(),
-            typedBody.type()
+            resumeType
         );
 
-        var typed = new TypedFunctionDefinitionNode(
+        var typed = new TypedTryHandlerMethodNode(
             untyped.name(),
-            typedSignature,
+            typedParams,
             typedBody,
             type,
             untyped.source()