Include effect params in constraint solving

2 files changed, 249 insertions, 35 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();
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
index 7821c68e..a1928bf8 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.typechecker;
 
 import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.sources.Sources;
+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.*;
@@ -11,12 +12,123 @@ import java.util.List;
 import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubEffect;
 import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubType;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.*;
 
 public class TypeCheckerIsSubTypeTests {
-    // == General ==
+    // == General effects ==
+
+    @Test
+    public void effectsAreSubeffectsOfThemselves() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var typeInfos = TypeInfos.EMPTY;
+
+        assertThat(isSubEffect(Effects.IO, Effects.IO, typeInfos), equalTo(true));
+        assertThat(
+            isSubEffect(
+                new UserDefinedEffect(namespaceName, "E"),
+                new UserDefinedEffect(namespaceName, "E"),
+                typeInfos
+            ),
+            equalTo(true)
+        );
+    }
+
+    @Test
+    public void unrelatedEffectsAreNotSubeffects() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var typeInfos = TypeInfos.EMPTY;
+
+        assertThat(
+            isSubEffect(
+                new UserDefinedEffect(namespaceName, "E1"),
+                new UserDefinedEffect(namespaceName, "E2"),
+                typeInfos
+            ),
+            equalTo(false)
+        );
+    }
+
+    // == None effect ==
+
+    @Test
+    public void noneEffectIsSubEffectOfAllEffects() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var typeInfos = TypeInfos.EMPTY;
+
+        assertThat(isSubEffect(Effects.NONE, Effects.IO, typeInfos), equalTo(true));
+        assertThat(
+            isSubEffect(
+                Effects.NONE,
+                new UserDefinedEffect(namespaceName, "E"),
+                typeInfos
+            ),
+            equalTo(true)
+        );
+    }
+
+    @Test
+    public void noneEffectIsSupereffectOfNoOtherEffects() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var typeInfos = TypeInfos.EMPTY;
+
+        assertThat(isSubEffect(Effects.IO, Effects.NONE, typeInfos), equalTo(false));
+        assertThat(
+            isSubEffect(
+                new UserDefinedEffect(namespaceName, "E"),
+                Effects.NONE,
+                typeInfos
+            ),
+            equalTo(false)
+        );
+    }
+
+    // == Effect params ==
+
+    @Test
+    public void whenSupereffectIsEffectArgThenEffectArgIsConstrained() {
+        var effectParam = EffectParam.forStructType(
+            Types.struct(NamespaceName.of("Example"), "Struct"),
+            "E"
+        );
+        var effectArg = new EffectArg(effectParam);
+
+        var constraints = TypeCheckerSubTyping.addSubeffectConstraint(
+            Effects.IO,
+            effectArg,
+            Sources.NULL,
+            TypeCheckerSubTyping.constraints(List.of(effectArg), TypeInfos.EMPTY)
+        );
+
+        assertThat(constraints.effectBounds().get(effectArg), isSequence(
+            equalTo(Effects.IO)
+        ));
+    }
+
+    @Test
+    public void whenSupereffectIsFreeEffectParamThenOtherEffectsAreNotSubtype() {
+        var effectParam = EffectParam.forStructType(
+            Types.struct(NamespaceName.of("Example"), "Struct"),
+            "E"
+        );
+
+        var error = assertThrows(
+            UnexpectedEffectError.class,
+            () -> TypeCheckerSubTyping.addSubeffectConstraint(
+                Effects.IO,
+                effectParam,
+                Sources.NULL,
+                TypeCheckerSubTyping.constraints(List.of(), TypeInfos.EMPTY)
+            )
+        );
+
+        assertThat(error.allowed(), equalTo(effectParam));
+        assertThat(error.actual(), equalTo(Effects.IO));
+    }
+
+    // == General types ==
 
     @Test
     public void typesAreSubTypesOfThemselves() {
@@ -586,6 +698,42 @@ public class TypeCheckerIsSubTypeTests {
     // == Constraint solving ==
 
     @Test
+    public void whenEffectParamHasNoBoundThenEffectParamIsResolvedToNoneEffect() {
+        var effectParam = EffectParam.forStructType(
+            Types.struct(NamespaceName.of("Example"), "Struct"),
+            "E"
+        );
+        var effectArg = new EffectArg(effectParam);
+
+        var constraints = TypeCheckerSubTyping.constraints(List.of(effectArg), TypeInfos.EMPTY);
+        var result = constraints.solve(Sources.NULL);
+
+        assertThat(result, equalTo(TypeSubstitutionMap.of(
+            Map.entry(effectParam, Effects.NONE)
+        )));
+    }
+
+    @Test
+    public void whenEffectParamHasLowerBoundsThenEffectParamIsResolvedToUnionOfBounds() {
+        var effect1 = new UserDefinedEffect(NamespaceName.of("Example"), "E1");
+        var effect2 = new UserDefinedEffect(NamespaceName.of("Example"), "E2");
+        var effectParam = EffectParam.forStructType(
+            Types.struct(NamespaceName.of("Example"), "Struct"),
+            "E"
+        );
+        var effectArg = new EffectArg(effectParam);
+
+        var constraints = TypeCheckerSubTyping.constraints(List.of(effectArg), TypeInfos.EMPTY)
+            .addLowerEffectBound(effectArg, effect1)
+            .addLowerEffectBound(effectArg, effect2);
+        var result = constraints.solve(Sources.NULL);
+
+        assertThat(result, equalTo(TypeSubstitutionMap.of(
+            Map.entry(effectParam, Effects.union(effect1, effect2))
+        )));
+    }
+
+    @Test
     public void whenTypeParamHasUpperBoundThenTypeParamIsResolvedToUpperBound() {
         var typeParam = TypeParam.forStructType(
             Types.struct(NamespaceName.of("Example"), "Struct"),
@@ -594,7 +742,7 @@ public class TypeCheckerIsSubTypeTests {
         var typeArg = new TypeArg(typeParam);
 
         var constraints = TypeCheckerSubTyping.constraints(List.of(typeArg), TypeInfos.EMPTY)
-            .addUpperBound(typeArg, Types.INT);
+            .addUpperTypeBound(typeArg, Types.INT);
         var result = constraints.solve(Sources.NULL);
 
         assertThat(result, equalTo(TypeSubstitutionMap.of(
@@ -611,7 +759,7 @@ public class TypeCheckerIsSubTypeTests {
         var typeArg = new TypeArg(typeParam);
 
         var constraints = TypeCheckerSubTyping.constraints(List.of(typeArg), TypeInfos.EMPTY)
-            .addLowerBound(typeArg, Types.INT);
+            .addLowerTypeBound(typeArg, Types.INT);
         var result = constraints.solve(Sources.NULL);
 
         assertThat(result, equalTo(TypeSubstitutionMap.of(
@@ -628,8 +776,8 @@ public class TypeCheckerIsSubTypeTests {
         var typeArg = new TypeArg(typeParam);
 
         var constraints = TypeCheckerSubTyping.constraints(List.of(typeArg), TypeInfos.EMPTY)
-            .addLowerBound(typeArg, Types.INT)
-            .addLowerBound(typeArg, Types.STRING);
+            .addLowerTypeBound(typeArg, Types.INT)
+            .addLowerTypeBound(typeArg, Types.STRING);
         var result = constraints.solve(Sources.NULL);
 
         assertThat(result, equalTo(TypeSubstitutionMap.of(
@@ -655,8 +803,8 @@ public class TypeCheckerIsSubTypeTests {
             )));
 
         var constraints = TypeCheckerSubTyping.constraints(List.of(typeArg), typeInfos)
-            .addLowerBound(typeArg, variantType1)
-            .addUpperBound(typeArg, sumType);
+            .addLowerTypeBound(typeArg, variantType1)
+            .addUpperTypeBound(typeArg, sumType);
         var result = constraints.solve(Sources.NULL);
 
         assertThat(result, equalTo(TypeSubstitutionMap.of(
@@ -682,8 +830,8 @@ public class TypeCheckerIsSubTypeTests {
             )));
 
         var constraints = TypeCheckerSubTyping.constraints(List.of(typeArg), typeInfos)
-            .addLowerBound(typeArg, sumType)
-            .addUpperBound(typeArg, variantType1);
+            .addLowerTypeBound(typeArg, sumType)
+            .addUpperTypeBound(typeArg, variantType1);
         var error = assertThrows(
             // TODO: more specific error
             IllegalStateException.class,