Type-check effect handler fields

2 files changed, 154 insertions, 73 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
index e86a2840..68541f64 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
@@ -6,34 +6,43 @@ import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.Effect;
 import org.shedlang.compiler.types.Effects;
 import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
 import java.util.Optional;
 
 public record PartialFunctionSignature(
-    List<Param> params,
+    List<Param> positionalParams,
+    List<Param> namedParams,
     Effect effect,
     Optional<Type> returnType
 ) {
     public record Param(String name, Type type, Source source) {}
 
+    public List<Param> params() {
+        return Lists.concat(positionalParams, namedParams);
+    }
+
     public static Builder builder() {
         return new Builder(
             P.vector(),
+            P.vector(),
             Effects.NONE,
             Optional.empty()
         );
     }
 
     public record Builder(
-        PVector<Param> params,
+        PVector<Param> positionalParams,
+        PVector<Param> namedParams,
         Effect effect,
         Optional<Type> returnType
     ) {
         public PartialFunctionSignature build() {
             return new PartialFunctionSignature(
-                params,
+                positionalParams,
+                namedParams,
                 effect,
                 returnType
             );
@@ -41,7 +50,8 @@ public record PartialFunctionSignature(
 
         public Builder addPositionalParam(String name, Type type) {
             return new Builder(
-                this.params.plus(new Param(name, type, Sources.NULL)),
+                this.positionalParams.plus(new Param(name, type, Sources.NULL)),
+                this.namedParams,
                 this.effect,
                 this.returnType
             );
@@ -49,7 +59,8 @@ public record PartialFunctionSignature(
 
         public Builder addNamedParam(String name, Type type) {
             return new Builder(
-                this.params.plus(new Param(name, type, Sources.NULL)),
+                this.positionalParams,
+                this.namedParams.plus(new Param(name, type, Sources.NULL)),
                 this.effect,
                 this.returnType
             );
@@ -57,7 +68,8 @@ public record PartialFunctionSignature(
 
         public Builder withEffect(Effect effect) {
             return new Builder(
-                this.params,
+                this.positionalParams,
+                this.namedParams,
                 effect,
                 this.returnType
             );
@@ -65,7 +77,8 @@ public record PartialFunctionSignature(
 
         public Builder withReturnType(Type returnType) {
             return new Builder(
-                this.params,
+                this.positionalParams,
+                this.namedParams,
                 this.effect,
                 Optional.of(returnType)
             );
@@ -73,7 +86,8 @@ public record PartialFunctionSignature(
 
         public Builder withoutReturnType() {
             return new Builder(
-                this.params,
+                this.positionalParams,
+                this.namedParams,
                 this.effect,
                 Optional.empty()
             );
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index a1139a98..03c60c93 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -701,15 +701,14 @@ public class TypeChecker {
             .map(typedParamNodeToParam)
             .toList();
 
-        var partialSignatureParams = Lists.concat(positionalParams, namedParams);
-
         var effect = functionSignatureToEffect(typedSignature);
 
         var returnType = functionSignatureToReturnType(typedSignature)
             .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(untyped.signature().source()));
 
         var partialSignature = new PartialFunctionSignature(
-            partialSignatureParams,
+            positionalParams,
+            namedParams,
             effect,
             Optional.of(returnType)
         );
@@ -766,55 +765,7 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
-        var functionTypeHint = Optionals.filterToType(FunctionType.class, hint);
-
-        var positionalParams = new ArrayList<PartialFunctionSignature.Param>();
-
-        for (var i = 0; i < typedSignature.params().positional().size(); i++) {
-            var paramIndex = i;
-            var paramNode = typedSignature.params().positional().get(paramIndex);
-            var paramHint = functionTypeHint.flatMap(h -> Lists.tryGet(h.positionalParams(), paramIndex));
-
-            var param = new PartialFunctionSignature.Param(
-                paramNode.name(),
-                paramNode.type()
-                    .map(x -> x.value())
-                    .or(() -> paramHint)
-                    .orElseThrow(() -> new FunctionSignatureMissingParamTypeError(paramNode.source())),
-                paramNode.source()
-            );
-            positionalParams.add(param);
-        }
-
-        var namedParams = new ArrayList<PartialFunctionSignature.Param>();
-
-        for (var namedParam : typedSignature.params().named()) {
-            var paramHint = functionTypeHint.flatMap(
-                h -> h.namedParams().stream()
-                    .filter(p -> p.name().equals(namedParam.name()))
-                    .map(p -> p.type())
-                    .findFirst()
-            );
-
-            namedParams.add(new PartialFunctionSignature.Param(
-                namedParam.name(),
-                namedParam.type()
-                    .map(x -> x.value())
-                    .or(() -> paramHint)
-                    .orElseThrow(() -> new FunctionSignatureMissingParamTypeError(namedParam.source())),
-                namedParam.source()
-            ));
-        }
-
-        var partialSignatureParams = Lists.concat(positionalParams, namedParams);
-
-        var effect = functionSignatureToEffect(typedSignature);
-
-        var partialSignature = new PartialFunctionSignature(
-            partialSignatureParams,
-            effect,
-            typedSignature.returnType().map(node -> node.value())
-        );
+        var partialSignature = inferFunctionPartialSignature(typedSignature, hint);
 
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
@@ -825,13 +776,13 @@ public class TypeChecker {
         context = typedBodyResult.context();
 
         var type = new FunctionType(
-            positionalParams.stream()
+            partialSignature.positionalParams().stream()
                 .map(x -> x.type())
                 .toList(),
-            namedParams.stream()
+            partialSignature.namedParams().stream()
                 .map(x -> new NamedParam(x.name(), x.type()))
                 .toList(),
-            effect,
+            partialSignature.effect(),
             typedBody.type()
         );
 
@@ -923,6 +874,60 @@ public class TypeChecker {
         );
     }
 
+    private static PartialFunctionSignature inferFunctionPartialSignature(
+        TypedFunctionSignatureNode typedSignature,
+        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++) {
+            var paramIndex = i;
+            var paramNode = typedSignature.params().positional().get(paramIndex);
+            var paramHint = functionTypeHint.flatMap(h -> Lists.tryGet(h.positionalParams(), paramIndex));
+
+            var param = new PartialFunctionSignature.Param(
+                paramNode.name(),
+                paramNode.type()
+                    .map(x -> x.value())
+                    .or(() -> paramHint)
+                    .orElseThrow(() -> new FunctionSignatureMissingParamTypeError(paramNode.source())),
+                paramNode.source()
+            );
+            positionalParams.add(param);
+        }
+
+        var namedParams = new ArrayList<PartialFunctionSignature.Param>();
+
+        for (var namedParam : typedSignature.params().named()) {
+            var paramHint = functionTypeHint.flatMap(
+                h -> h.namedParams().stream()
+                    .filter(p -> p.name().equals(namedParam.name()))
+                    .map(p -> p.type())
+                    .findFirst()
+            );
+
+            namedParams.add(new PartialFunctionSignature.Param(
+                namedParam.name(),
+                namedParam.type()
+                    .map(x -> x.value())
+                    .or(() -> paramHint)
+                    .orElseThrow(() -> new FunctionSignatureMissingParamTypeError(namedParam.source())),
+                namedParam.source()
+            ));
+        }
+
+        var effect = functionSignatureToEffect(typedSignature);
+
+        return new PartialFunctionSignature(
+            positionalParams,
+            namedParams,
+            effect,
+            typedSignature.returnType().map(node -> node.value())
+        );
+    }
+
     private static Effect functionSignatureToEffect(TypedFunctionSignatureNode signature) {
         return signature.effect().map(e -> e.value()).orElse(Effects.NONE);
     }
@@ -1617,27 +1622,89 @@ public class TypeChecker {
         context = context.exit();
 
         var effectInfo = context.userDefinedEffectInfo(handledEffect, typedHandledEffect.source());
-        // TODO: Don't create unnecessary list
-        for (var effectField : effectInfo.fields().stream().toList()) {
-            throw new MissingEffectHandlerFieldError(effectField.name(), untyped.source());
-        }
+        var handledFieldNames = new HashSet<String>();
+
+        var typedHandlerFields = new ArrayList<TypedFunctionDefinitionNode>();
+        for (var untypedHandlerField : untyped.handlers()) {
+            handledFieldNames.add(untypedHandlerField.name());
+
+            var effectFieldType = effectInfo.fields()
+                .fieldType(untypedHandlerField.name())
+                .orElseThrow(() -> new NoSuchFieldError(
+                    Types.effectType(handledEffect),
+                    untypedHandlerField.name(),
+                    untypedHandlerField.source()
+                ));
+
+            var typedHandlerFieldResult = typeCheckTryHandlerField(untypedHandlerField, effectFieldType, context);
+            var typedHandlerField = typedHandlerFieldResult.value();
+            context = typedHandlerFieldResult.context();
 
-        for (var handlerField : untyped.handlers()) {
-            throw new NoSuchFieldError(
-                Types.effectType(handledEffect),
-                handlerField.name(),
-                handlerField.source()
+            checkIsSubType(
+                typedHandlerField.type(),
+                effectFieldType,
+                untypedHandlerField.source(),
+                context
             );
+
+            typedHandlerFields.add(typedHandlerField);
+        }
+
+        // TODO: Don't create unnecessary list
+        for (var effectField : effectInfo.fields().stream().toList()) {
+            if (!handledFieldNames.contains(effectField.name())) {
+                throw new MissingEffectHandlerFieldError(effectField.name(), untyped.source());
+            }
         }
 
         return context.toResult(new TypedTryNode(
             typedBody,
             typedHandledEffect,
-            P.vector(),
+            typedHandlerFields,
             untyped.source()
         ));
     }
 
+    private static TypeCheckerResult<TypedFunctionDefinitionNode> typeCheckTryHandlerField(
+        UntypedFunctionDefinitionNode untyped,
+        FunctionType effectFieldType,
+        TypeCheckerContext context
+    ) {
+        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
+
+        var partialSignature = inferFunctionPartialSignature(typedSignature, Optional.of(effectFieldType));
+
+        // TODO: don't allow returns, allow resumes
+        var typedBodyResult = typeCheckFunctionBody(
+            untyped.body(),
+            partialSignature,
+            context
+        );
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
+        var type = new FunctionType(
+            partialSignature.positionalParams().stream()
+                .map(x -> x.type())
+                .toList(),
+            partialSignature.namedParams().stream()
+                .map(x -> new NamedParam(x.name(), x.type()))
+                .toList(),
+            partialSignature.effect(),
+            typedBody.type()
+        );
+
+        var typed = new TypedFunctionDefinitionNode(
+            untyped.name(),
+            typedSignature,
+            typedBody,
+            type,
+            untyped.source()
+        );
+
+        return context.toResult(typed);
+    }
+
     private static TypeCheckerResult<TypedExpressionNode> typeCheckUnitLiteral(
         UntypedUnitLiteralNode untyped,
         TypeCheckerContext context