Include effect params in constraint solving

1 files changed, 92 insertions, 26 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index f1145d5a..935d1078 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -3,6 +3,7 @@ package org.shedlang.compiler.typechecker;
 import org.pcollections.OrderedPMap;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.typechecker.errors.UnexpectedEffectError;
 import org.shedlang.compiler.typechecker.errors.UnexpectedTypeError;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
@@ -18,7 +19,7 @@ public class TypeCheckerSubTyping {
     private TypeCheckerSubTyping() {}
 
     public static boolean isSubType(Type subtype, Type supertype, TypeInfos typeInfos) {
-        var constraints = new Constraints(P.orderedPMap(), typeInfos);
+        var constraints = constraints(typeInfos);
         return tryAddSubtypeConstraint(subtype, supertype, constraints)
             .isPresent();
     }
@@ -41,12 +42,18 @@ public class TypeCheckerSubTyping {
         }
     }
 
+    private static Constraints constraints(TypeInfos typeInfos) {
+        return new Constraints(P.orderedPMap(), P.orderedPMap(), typeInfos);
+    }
+
     static Constraints constraints(List<TypeLevelArg> typeLevelArgs, TypeInfos typeInfos) {
+        var effectBounds = P.<EffectArg, PVector<Effect>>orderedPMap();
         var typeBounds = P.<TypeArg, TypeBounds>orderedPMap();
 
         for (var typeLevelArg : typeLevelArgs) {
             switch (typeLevelArg) {
                 case EffectArg effectArg -> {
+                    effectBounds = effectBounds.plus(effectArg, P.vector());
                 }
                 case TypeArg typeArg -> {
                     typeBounds = typeBounds.plus(typeArg, TypeBounds.EMPTY);
@@ -54,10 +61,11 @@ public class TypeCheckerSubTyping {
             }
         }
 
-        return new Constraints(typeBounds, typeInfos);
+        return new Constraints(effectBounds, typeBounds, typeInfos);
     }
 
     record Constraints(
+        OrderedPMap<EffectArg, PVector<Effect>> effectBounds,
         OrderedPMap<TypeArg, TypeBounds> typeBounds,
         TypeInfos typeInfos
     ) {
@@ -65,20 +73,36 @@ public class TypeCheckerSubTyping {
             return this.typeBounds.containsKey(typeArg);
         }
 
-        public Constraints addLowerBound(TypeArg typeArg, Type lowerBound) {
-            return updateBounds(typeArg, bounds -> bounds.addLowerBound(lowerBound));
+        public boolean isBoundEffectArg(EffectArg effectArg) {
+            return this.effectBounds.containsKey(effectArg);
+        }
+
+        public Constraints addLowerEffectBound(EffectArg effectArg, Effect subeffect) {
+            return new Constraints(
+                this.effectBounds.plus(
+                    effectArg,
+                    this.effectBounds.get(effectArg).plus(subeffect)
+                ),
+                this.typeBounds,
+                this.typeInfos
+            );
+        }
+
+        public Constraints addLowerTypeBound(TypeArg typeArg, Type lowerBound) {
+            return updateTypeBounds(typeArg, bounds -> bounds.addLowerBound(lowerBound));
         }
 
-        public Constraints addUpperBound(TypeArg typeArg, Type upperBound) {
-            return updateBounds(typeArg, bounds -> bounds.addUpperBound(upperBound));
+        public Constraints addUpperTypeBound(TypeArg typeArg, Type upperBound) {
+            return updateTypeBounds(typeArg, bounds -> bounds.addUpperBound(upperBound));
         }
 
-        private Constraints updateBounds(
+        private Constraints updateTypeBounds(
             TypeArg typeArg,
             Function<TypeBounds, TypeBounds> update
         ) {
             var bounds = update.apply(this.typeBounds.get(typeArg));
             return new Constraints(
+                this.effectBounds,
                 this.typeBounds.plus(typeArg, bounds),
                 this.typeInfos
             );
@@ -90,17 +114,24 @@ public class TypeCheckerSubTyping {
 
             var map = new HashMap<TypeLevelParam, TypeLevelValue>();
 
+            for (var effectBoundsEntry : this.effectBounds.entrySet()) {
+                var effectArg = effectBoundsEntry.getKey();
+                var effectBounds = effectBoundsEntry.getValue();
+                var inferredEffect = effectBounds.stream().reduce(Effects.NONE, Effects::union);
+                map.put(effectArg.param(), inferredEffect);
+            }
+
             for (var typeBoundsEntry : this.typeBounds.entrySet()) {
                 var typeArg = typeBoundsEntry.getKey();
                 var typeBounds = typeBoundsEntry.getValue();
-                var inferredType = solveBounds(typeBounds);
+                var inferredType = solveTypeBounds(typeBounds);
                 map.put(typeArg.param(), inferredType);
             }
 
             return new TypeSubstitutionMap(map);
         }
 
-        private Type solveBounds(TypeBounds typeParamBounds) {
+        private Type solveTypeBounds(TypeBounds typeParamBounds) {
             var lowerBound = Types.BOTTOM;
             for (var lower : typeParamBounds.lower) {
                 lowerBound = Types.commonSupertype(
@@ -128,6 +159,52 @@ public class TypeCheckerSubTyping {
         }
     }
 
+    static Constraints addSubeffectConstraint(
+        Effect subeffect,
+        Effect supereffect,
+        Source source,
+        Constraints constraints
+    ) {
+        var result = tryAddSubeffectConstraint(subeffect, supereffect, constraints);
+        if (result.isEmpty()) {
+            throw new UnexpectedEffectError(
+                supereffect,
+                subeffect,
+                source
+            );
+        } else {
+            return result.get();
+        }
+    }
+
+    private static Optional<Constraints> tryAddSubeffectConstraint(
+        Effect subeffect,
+        Effect supereffect,
+        Constraints constraints
+    ) {
+        if (subeffect == Effects.NONE) {
+            return Optional.of(constraints);
+        }
+
+        if (subeffect.equals(supereffect)) {
+            return Optional.of(constraints);
+        }
+
+        if (
+            supereffect instanceof EffectArg supereffectEffectArg &&
+            constraints.isBoundEffectArg(supereffectEffectArg)
+        ) {
+            return Optional.of(constraints.addLowerEffectBound(supereffectEffectArg, subeffect));
+        }
+
+        if (supereffect instanceof EffectUnion superEffectUnion) {
+            return tryAddSubeffectConstraint(subeffect, superEffectUnion.left(), constraints)
+                .or(() -> tryAddSubeffectConstraint(subeffect, superEffectUnion.right(), constraints));
+        }
+
+        return Optional.empty();
+    }
+
     static Constraints addSubtypeConstraint(
         Type subtype,
         Type supertype,
@@ -167,7 +244,7 @@ public class TypeCheckerSubTyping {
             supertype instanceof TypeArg supertypeTypeArg &&
             constraints.isBoundTypeArg(supertypeTypeArg)
         ) {
-            constraints = constraints.addLowerBound(supertypeTypeArg, subtype);
+            constraints = constraints.addLowerTypeBound(supertypeTypeArg, subtype);
 
             return Optional.of(constraints);
         }
@@ -176,7 +253,7 @@ public class TypeCheckerSubTyping {
             subtype instanceof TypeArg subtypeTypeArg &&
             constraints.isBoundTypeArg(subtypeTypeArg)
         ) {
-            constraints = constraints.addUpperBound(subtypeTypeArg, supertype);
+            constraints = constraints.addUpperTypeBound(subtypeTypeArg, supertype);
 
             return Optional.of(constraints);
         }
@@ -314,20 +391,9 @@ public class TypeCheckerSubTyping {
         return Optional.empty();
     }
 
-    static boolean isSubEffect(Effect subEffect, Effect superEffect, TypeInfos typeInfos) {
-        if (subEffect == Effects.NONE) {
-            return true;
-        }
-
-        if (subEffect.equals(superEffect)) {
-            return true;
-        }
-
-        if (superEffect instanceof EffectUnion superEffectUnion) {
-            return isSubEffect(subEffect, superEffectUnion.left(), typeInfos) ||
-                isSubEffect(subEffect, superEffectUnion.right(), typeInfos);
-        }
-
-        return false;
+    static boolean isSubEffect(Effect subeffect, Effect supereffect, TypeInfos typeInfos) {
+        var constraints = constraints(typeInfos);
+        return tryAddSubeffectConstraint(subeffect, supereffect, constraints)
+            .isPresent();
     }
 }