Support subtyping for constructed sum types

3 files changed, 124 insertions, 0 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index dcaa9226..ffd2ac01 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -481,6 +481,31 @@ public class TypeCheckerContext {
             .orElseThrow(() -> new MissingSumTypeInfoError(type, source));
     }
 
+    public SumTypeInfo sumTypeInfo(ConstructedType<SumType> type, Source source) {
+        var typeInfo = sumTypeInfo(type.genericType(), source);
+
+        var typeSubstitutionMap = type.typeSubstitutionMap();
+
+        var variants = typeInfo.variants().stream()
+            .map(variant -> switch (variant) {
+                case SumVariant.ConstructedStruct constructedStruct ->
+                    // TODO: overload to avoid cast?
+                    SumVariant.constructedStruct((ConstructedType<StructType>) TypeSubstitution.substitute(
+                        constructedStruct.type(),
+                        typeSubstitutionMap
+                    ));
+
+                case SumVariant.Struct struct ->
+                    SumVariant.struct((StructType) TypeSubstitution.substitute(
+                        struct.type(),
+                        typeSubstitutionMap
+                    ));
+            })
+            .collect(PCollectors.toOrderedSet());
+
+        return new SumTypeInfo(variants);
+    }
+
     public TypeCheckerContext addUserDefinedEffectInfo(
         UserDefinedEffect effect,
         UserDefinedEffectInfo effectInfo
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index 2c69ac14..4aab5370 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -29,6 +29,16 @@ public class TypeCheckerSubTyping {
         }
 
         if (
+            supertype instanceof ConstructedType<?> supertypeConstructedType &&
+            supertypeConstructedType.genericType() instanceof SumType
+        ) {
+            // TODO: merge with logic for plain sum types
+            var superTypeInfo = context.sumTypeInfo((ConstructedType<SumType>) supertypeConstructedType, Sources.NULL);
+            return superTypeInfo.variants().stream()
+                .anyMatch(supertypeVariant -> isSubType(subtype, supertypeVariant.toType(), context));
+        }
+
+        if (
             supertype instanceof FunctionType supertypeFunction &&
             subtype instanceof FunctionType subtypeFunction
         ) {
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
index 568256c2..ca3008f3 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
@@ -4,6 +4,8 @@ import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 
+import java.util.List;
+
 import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubType;
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.equalTo;
@@ -138,6 +140,93 @@ public class TypeCheckerIsSubTypeTests {
     }
 
     @Test
+    public void variantWithoutTypeLevelArgIsSubTypeOfConstructedSumType() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var sumType = Types.sum(namespaceName, "Sum");
+        var sumTypeConstructor = Types.typeConstructor(
+            List.of(TypeParam.forSumType(sumType, "T")),
+            sumType
+        );
+        var constructedSumType = Types.constructedType(
+            sumTypeConstructor,
+            List.of(Types.STRING)
+        );
+        var variantType1 = Types.struct(namespaceName, "Variant1");
+        var variantType2 = Types.struct(namespaceName, "Variant2");
+        var context = TypeCheckerContext.EMPTY
+            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(variantType1),
+                SumVariant.struct(variantType2)
+            )));
+
+        assertThat(isSubType(variantType1, constructedSumType, context), equalTo(true));
+        assertThat(isSubType(variantType2, constructedSumType, context), equalTo(true));
+    }
+
+    @Test
+    public void variantWithTypeLevelArgIsSubTypeOfConstructedSumType() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var sumType = Types.sum(namespaceName, "Sum");
+        var sumTypeConstructor = Types.typeConstructor(
+            List.of(TypeParam.forSumType(sumType, "T")),
+            sumType
+        );
+        var constructedSumType = Types.constructedType(
+            sumTypeConstructor,
+            List.of(Types.STRING)
+        );
+        var variantType1 = Types.struct(namespaceName, "Variant1");
+        var variantTypeConstructor1 = Types.typeConstructor(
+            List.of(TypeParam.forStructType(variantType1, "T")),
+            variantType1
+        );
+        var constructedVariantType1 = Types.constructedType(
+            variantTypeConstructor1,
+            List.of(Types.STRING)
+        );
+        var variantType2 = Types.struct(namespaceName, "Variant2");
+        var context = TypeCheckerContext.EMPTY
+            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(
+                SumVariant.constructedStruct(constructedVariantType1),
+                SumVariant.struct(variantType2)
+            )));
+
+        assertThat(isSubType(constructedVariantType1, constructedSumType, context), equalTo(true));
+        assertThat(
+            isSubType(
+                Types.constructedType(variantTypeConstructor1, List.of(Types.INT)),
+                constructedSumType,
+                context
+            ),
+            equalTo(false)
+        );
+    }
+
+    @Test
+    public void variantIsNotSuperTypeOfConstructedSumType() {
+        var namespaceName = NamespaceName.of("Example", "Project");
+        var sumType = Types.sum(namespaceName, "Sum");
+        var sumTypeConstructor = Types.typeConstructor(
+            List.of(TypeParam.forSumType(sumType, "T")),
+            sumType
+        );
+        var constructedSumType = Types.constructedType(
+            sumTypeConstructor,
+            List.of(Types.STRING)
+        );
+        var variantType1 = Types.struct(namespaceName, "Variant1");
+        var variantType2 = Types.struct(namespaceName, "Variant2");
+        var context = TypeCheckerContext.EMPTY
+            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(variantType1),
+                SumVariant.struct(variantType2)
+            )));
+
+        assertThat(isSubType(constructedSumType, variantType1, context), equalTo(false));
+        assertThat(isSubType(constructedSumType, variantType2, context), equalTo(false));
+    }
+
+    @Test
     public void whenFunctionsHaveDifferentPositionalParamCountsThenTheyAreNotSubtypes() {
         var functionType1 = FunctionType.builder()
             .addPositionalParam(Types.STRING)