Add type level params to environment for sum variants

3 files changed, 102 insertions, 12 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
index d6bd14be..568256c2 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIsSubTypeTests.java
@@ -112,7 +112,10 @@ public class TypeCheckerIsSubTypeTests {
         var variantType1 = Types.struct(namespaceName, "Variant1");
         var variantType2 = Types.struct(namespaceName, "Variant2");
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(variantType1, variantType2)));
+            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(variantType1),
+                SumVariant.struct(variantType2)
+            )));
 
         assertThat(isSubType(variantType1, sumType, context), equalTo(true));
         assertThat(isSubType(variantType2, sumType, context), equalTo(true));
@@ -125,7 +128,10 @@ public class TypeCheckerIsSubTypeTests {
         var variantType1 = Types.struct(namespaceName, "Variant1");
         var variantType2 = Types.struct(namespaceName, "Variant2");
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(variantType1, variantType2)));
+            .addSumTypeInfo(sumType, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(variantType1),
+                SumVariant.struct(variantType2)
+            )));
 
         assertThat(isSubType(sumType, variantType1, context), equalTo(false));
         assertThat(isSubType(sumType, variantType2, context), equalTo(false));
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java
index a028b21f..33d254cd 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java
@@ -6,6 +6,8 @@ import org.shedlang.compiler.ast.untyped.UntypedSumTypeDefinitionNode;
 import org.shedlang.compiler.typechecker.errors.CannotDefineSumTypeOutsideOfNamespaceError;
 import org.shedlang.compiler.typechecker.errors.VariantTypeIsNotStructError;
 import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.types.SumVariant;
+import org.shedlang.compiler.types.TypeParam;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
 
@@ -155,8 +157,8 @@ public class TypeCheckerSumTypeDefinitionTests {
         assertThat(
             sumTypeInfo,
             has("variants", x -> x.variants(), isSequence(
-                equalTo(variantType1),
-                equalTo(variantType2)
+                equalTo(SumVariant.struct(variantType1)),
+                equalTo(SumVariant.struct(variantType2))
             ))
         );
     }
@@ -179,4 +181,67 @@ public class TypeCheckerSumTypeDefinitionTests {
 
         assertThat(error.variantType(), equalTo(Types.STRING));
     }
+
+    @Test
+    public void typeParamsCanBeUsedInVariants() {
+        var untyped = UntypedSumTypeDefinitionNode.builder()
+            .withName("ExampleSumType")
+            .withTypeLevelParams(List.of(
+                Untyped.typeLevelParam("T1"),
+                Untyped.typeLevelParam("T2")
+            ))
+            .addVariant(Untyped.constructedType(
+                Untyped.typeLevelReference("A"),
+                List.of(
+                    Untyped.typeLevelReference("T1")
+                )
+            ))
+            .addVariant(Untyped.constructedType(
+                Untyped.typeLevelReference("B"),
+                List.of(
+                    Untyped.typeLevelReference("T2")
+                )
+            ))
+            .build();
+        var variantType1 = Types.struct(NAMESPACE_NAME, "A");
+        var variantTypeConstructor1 = Types.typeConstructor(
+            List.of(TypeParam.forStructType(variantType1, "T")),
+            variantType1
+        );
+        var variantType2 = Types.struct(NAMESPACE_NAME, "B");
+        var variantTypeConstructor2 = Types.typeConstructor(
+            List.of(TypeParam.forStructType(variantType2, "T")),
+            variantType2
+        );
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map(
+                Map.entry("A", Types.typeConstructorType(variantTypeConstructor1)),
+                Map.entry("B", Types.typeConstructorType(variantTypeConstructor2))
+            ));
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        var sumTypeInfo = typed.context().toTypeInfos()
+            .sumTypeInfo(Types.sum(NAMESPACE_NAME, "ExampleSumType"))
+            .orElseThrow();
+        assertThat(
+            sumTypeInfo,
+            has("variants", x -> x.variants(), isSequence(
+                equalTo(SumVariant.constructedStruct(Types.constructedType(
+                    variantTypeConstructor1,
+                    List.of(TypeParam.forSumType(
+                        Types.sum(NAMESPACE_NAME, "ExampleSumType"),
+                        "T1"
+                    ))
+                ))),
+                equalTo(SumVariant.constructedStruct(Types.constructedType(
+                    variantTypeConstructor2,
+                    List.of(TypeParam.forSumType(
+                        Types.sum(NAMESPACE_NAME, "ExampleSumType"),
+                        "T2"
+                    ))
+                )))
+            ))
+        );
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
index 252d8471..a3fc4bab 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
@@ -57,7 +57,10 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1, VARIANT_TYPE_2)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1),
+                SumVariant.struct(VARIANT_TYPE_2)
+            )))
             .enterGlobal(P.map(
                 Map.entry("x", SUM_TYPE),
                 Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1)),
@@ -69,7 +72,7 @@ public class TypeCheckerWhenTests {
             () -> typeCheckExpression(untyped, context)
         );
 
-        assertThat(error.unhandledVariant(), equalTo(VARIANT_TYPE_2));
+        assertThat(error.unhandledVariant(), equalTo(SumVariant.struct(VARIANT_TYPE_2)));
     }
 
     @Test
@@ -93,7 +96,9 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1)
+            )))
             .enterGlobal(P.map(
                 Map.entry("x", SUM_TYPE),
                 Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1)),
@@ -128,7 +133,9 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1)
+            )))
             .enterGlobal(P.map(
                 Map.entry("x", SUM_TYPE),
                 Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1))
@@ -162,7 +169,10 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1, VARIANT_TYPE_2)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1),
+                SumVariant.struct(VARIANT_TYPE_2)
+            )))
             .enterGlobal(P.map(
                 Map.entry("x", SUM_TYPE),
                 Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1)),
@@ -210,7 +220,10 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1, VARIANT_TYPE_2)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1),
+                SumVariant.struct(VARIANT_TYPE_2)
+            )))
             .enterGlobal(P.map(
                 Map.entry("x", SUM_TYPE),
                 Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1)),
@@ -242,7 +255,10 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1, VARIANT_TYPE_2)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1),
+                SumVariant.struct(VARIANT_TYPE_2)
+            )))
             .enterGlobal(P.map(
                 Map.entry("x", SUM_TYPE),
                 Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1)),
@@ -280,7 +296,10 @@ public class TypeCheckerWhenTests {
             )
         );
         var context = TypeCheckerContext.EMPTY
-            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1, VARIANT_TYPE_2)))
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                SumVariant.struct(VARIANT_TYPE_1),
+                SumVariant.struct(VARIANT_TYPE_2)
+            )))
             .addStructTypeInfo(
                 VARIANT_TYPE_1,
                 StructTypeInfo.builder()