Support inferring type args in when cases

1 files changed, 73 insertions, 0 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
index ecf2a137..e41bc2fa 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
@@ -264,6 +264,79 @@ public class TypeCheckerWhenTests {
     }
 
     @Test
+    public void variantCanBeGenericStructTypeWithInferredTypeArgs() {
+        // TODO: do we need to consider sums with multiple variants referring to the same raw type?
+        var variantTypeConstructor1 = Types.typeConstructor(
+            List.of(TypeParam.forStructType(VARIANT_TYPE_1, "T")),
+            VARIANT_TYPE_1
+        );
+        var constructedVariantType1 = Types.constructedStructType(
+            variantTypeConstructor1,
+            List.of(Types.STRING)
+        );
+        var untyped = Untyped.when(
+            Untyped.reference("x"),
+            List.of(
+                Untyped.whenCase(
+                    Untyped.typeLevelReference(VARIANT_TYPE_1.name()),
+                    Untyped.block(List.of(
+                        Untyped.blockReturn(
+                            Untyped.fieldAccess(
+                                Untyped.reference("x"),
+                                "value"
+                            )
+                        )
+                    ))
+                ),
+                Untyped.whenCase(
+                    Untyped.typeLevelReference(VARIANT_TYPE_2.name()),
+                    Untyped.block(List.of(
+                        Untyped.blockReturn(Untyped.stringLiteral("b"))
+                    ))
+                )
+            )
+        );
+        var context = TypeCheckerContext.EMPTY
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(
+                constructedVariantType1,
+                VARIANT_TYPE_2
+            )))
+            .addStructTypeInfo(
+                VARIANT_TYPE_1,
+                new StructTypeInfo(
+                    FieldsInfo.builder()
+                        .addField("value", TypeParam.forStructType(VARIANT_TYPE_1, "T"))
+                        .build()
+                )
+            )
+            .enterGlobal(P.map(
+                Map.entry("x", SUM_TYPE),
+                Map.entry(VARIANT_TYPE_1.name(), Types.typeConstructorType(variantTypeConstructor1)),
+                Map.entry(VARIANT_TYPE_2.name(), Types.metaType(VARIANT_TYPE_2))
+            ));
+
+        var result = typeCheckExpression(untyped, context);
+
+        assertThat(result.value(), isTypedWhenNode(
+            isTypedReferenceNode("x", SUM_TYPE),
+            isSequence(
+                isTypedWhenCaseNode(
+                    equalTo(constructedVariantType1),
+                    isTypedBlockNode(isSequence(
+                        isTypedBlockReturnNode(isTypedExpressionNode(Types.STRING))
+                    ))
+                ),
+                isTypedWhenCaseNode(
+                    equalTo(VARIANT_TYPE_2),
+                    isTypedBlockNode(isSequence(
+                        isTypedBlockReturnNode(isTypedStringLiteralNode("b"))
+                    ))
+                )
+            )
+        ));
+    }
+
+    @Test
     public void whenCaseHasConstructedTypeWithWrongTypeArgThenErrorIsThrown() {
         var variantTypeConstructor1 = Types.typeConstructor(
             List.of(TypeParam.forStructType(VARIANT_TYPE_1, "T")),