Support inferring type args in when cases

1 files changed, 37 insertions, 4 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 08600f43..187442a8 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -2101,18 +2101,20 @@ public class TypeChecker {
             ? Optional.of(expressionReference.name())
             : Optional.<String>empty();
 
+        var variantTypes = expressionTypeInfo.variants();
+
         var typedCasesResult = context.mapWithContext(
             untyped.cases(),
             (untypedWhenCase, whenCaseContext) -> typeCheckWhenCase(
                 untypedWhenCase,
                 variableName,
+                variantTypes,
                 whenCaseContext
             )
         );
         var typedCases = typedCasesResult.value();
         context = typedCasesResult.context();
 
-        var variantTypes = expressionTypeInfo.variants();
         var handledVariantTypes = new HashSet<Type>();
 
         for (var typedCase : typedCases) {
@@ -2149,15 +2151,46 @@ public class TypeChecker {
     private static TypeCheckerResult<TypedWhenCaseNode> typeCheckWhenCase(
         UntypedWhenCaseNode untyped,
         Optional<String> variableName,
+        OrderedPSet<StructType> variantTypes,
         TypeCheckerContext context
     ) {
         var typedCaseType = typeCheckTypeLevelExpression(
-            TypeSets.STRUCT_META_TYPE,
             untyped.caseType(),
             context
         );
 
-        var variant = typedCaseType.value();
+        var variant = switch (typedCaseType.value()) {
+            case StructType variantStructType ->
+                variantStructType;
+
+            case TypeConstructor<?> variantTypeConstructor when
+                variantTypeConstructor.genericType() instanceof StructType genericStructType -> {
+
+                var matchingVariantTypes = variantTypes.stream()
+                    .filter(
+                        variantType -> variantType instanceof ConstructedType<?> variantTypeConstructed &&
+                            variantTypeConstructed.genericType().equals(genericStructType)
+                    )
+                    .toList();
+
+                if (matchingVariantTypes.size() == 1) {
+                    yield matchingVariantTypes.get(0);
+                } else {
+                    throw new UnexpectedTypeError(
+                        TypeSets.STRUCT_META_TYPE,
+                        Types.typeLevelValueType(typedCaseType.value()),
+                        typedCaseType.source()
+                    );
+                }
+            }
+
+            default ->
+                throw new UnexpectedTypeError(
+                    TypeSets.STRUCT_META_TYPE,
+                    Types.typeLevelValueType(typedCaseType.value()),
+                    typedCaseType.source()
+                );
+        };
 
         var declaredVariables = DeclaredVariables.forFunctionStatements(
             untyped.body().statements(),
@@ -2168,7 +2201,7 @@ public class TypeChecker {
         context = context.enterBlock(bodyEnvironment.build());
 
         if (variableName.isPresent()) {
-            context = context.updateType(variableName.get(), typedCaseType.value(), untyped.source());
+            context = context.updateType(variableName.get(), variant, untyped.source());
         }
 
         var typedBodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.toBoundVariables(), context);