Discriminate contextual types using shorthand properties (#55151)

Author: Andarist

src/compiler/checker.ts

@@ -22785,7 +22785,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
             findMostOverlappyType(source, target);
     }
 
-    function discriminateTypeByDiscriminableItems(target: UnionType, discriminators: [() => Type, __String][], related: (source: Type, target: Type) => boolean | Ternary) {
+    function discriminateTypeByDiscriminableItems(target: UnionType, discriminators: (readonly [() => Type, __String])[], related: (source: Type, target: Type) => boolean | Ternary) {
         const types = target.types;
         const include: Ternary[] = types.map(t => t.flags & TypeFlags.Primitive ? Ternary.False : Ternary.True);
         for (const [getDiscriminatingType, propertyName] of discriminators) {
@@ -29438,12 +29438,23 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
         return getMatchingUnionConstituentForObjectLiteral(contextualType, node) || discriminateTypeByDiscriminableItems(contextualType,
             concatenate(
                 map(
-                    filter(node.properties, p => !!p.symbol && p.kind === SyntaxKind.PropertyAssignment && isPossiblyDiscriminantValue(p.initializer) && isDiscriminantProperty(contextualType, p.symbol.escapedName)),
-                    prop => ([() => getContextFreeTypeOfExpression((prop as PropertyAssignment).initializer), prop.symbol.escapedName] as [() => Type, __String])
+                    filter(node.properties, (p): p is PropertyAssignment | ShorthandPropertyAssignment => {
+                        if (!p.symbol) {
+                            return false;
+                        }
+                        if (p.kind === SyntaxKind.PropertyAssignment) {
+                            return isPossiblyDiscriminantValue(p.initializer) && isDiscriminantProperty(contextualType, p.symbol.escapedName);
+                        }
+                        if (p.kind === SyntaxKind.ShorthandPropertyAssignment) {
+                            return isDiscriminantProperty(contextualType, p.symbol.escapedName);
+                        }
+                        return false;
+                    }),
+                    prop => ([() => getContextFreeTypeOfExpression(prop.kind === SyntaxKind.PropertyAssignment ? prop.initializer : prop.name), prop.symbol.escapedName] as const)
                 ),
                 map(
                     filter(getPropertiesOfType(contextualType), s => !!(s.flags & SymbolFlags.Optional) && !!node?.symbol?.members && !node.symbol.members.has(s.escapedName) && isDiscriminantProperty(contextualType, s.escapedName)),
-                    s => [() => undefinedType, s.escapedName] as [() => Type, __String]
+                    s => [() => undefinedType, s.escapedName] as const
                 )
             ),
             isTypeAssignableTo
@@ -29456,7 +29467,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
             concatenate(
                 map(
                     filter(node.properties, p => !!p.symbol && p.kind === SyntaxKind.JsxAttribute && isDiscriminantProperty(contextualType, p.symbol.escapedName) && (!p.initializer || isPossiblyDiscriminantValue(p.initializer))),
-                    prop => ([!(prop as JsxAttribute).initializer ? (() => trueType) : (() => getContextFreeTypeOfExpression((prop as JsxAttribute).initializer!)), prop.symbol.escapedName] as [() => Type, __String])
+                    prop => ([!(prop as JsxAttribute).initializer ? (() => trueType) : (() => getContextFreeTypeOfExpression((prop as JsxAttribute).initializer!)), prop.symbol.escapedName] as const)
                 ),
                 map(
                     filter(getPropertiesOfType(contextualType), s => {
@@ -29469,7 +29480,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                         }
                         return !node.symbol.members.has(s.escapedName) && isDiscriminantProperty(contextualType, s.escapedName);
                     }),
-                    s => [() => undefinedType, s.escapedName] as [() => Type, __String]
+                    s => [() => undefinedType, s.escapedName] as const
                 )
             ),
             isTypeAssignableTo

tests/baselines/reference/contextuallyTypedByDiscriminableUnion.js

@@ -25,6 +25,14 @@ invoke({
         return +a;
     }
 });
+
+const kind = "a"
+invoke({
+    kind,
+    method(a) {
+        return +a;
+    }
+})
 
 
 //// [contextuallyTypedByDiscriminableUnion.js]
@@ -42,3 +50,10 @@ invoke({
         return +a;
     }
 });
+var kind = "a";
+invoke({
+    kind: kind,
+    method: function (a) {
+        return +a;
+    }
+});

tests/baselines/reference/contextuallyTypedByDiscriminableUnion.symbols

@@ -60,3 +60,21 @@ invoke({
     }
 });
 
+const kind = "a"
+>kind : Symbol(kind, Decl(contextuallyTypedByDiscriminableUnion.ts, 25, 5))
+
+invoke({
+>invoke : Symbol(invoke, Decl(contextuallyTypedByDiscriminableUnion.ts, 6, 2))
+
+    kind,
+>kind : Symbol(kind, Decl(contextuallyTypedByDiscriminableUnion.ts, 26, 8))
+
+    method(a) {
+>method : Symbol(method, Decl(contextuallyTypedByDiscriminableUnion.ts, 27, 9))
+>a : Symbol(a, Decl(contextuallyTypedByDiscriminableUnion.ts, 28, 11))
+
+        return +a;
+>a : Symbol(a, Decl(contextuallyTypedByDiscriminableUnion.ts, 28, 11))
+    }
+})
+

tests/baselines/reference/contextuallyTypedByDiscriminableUnion.types

@@ -69,3 +69,25 @@ invoke({
     }
 });
 
+const kind = "a"
+>kind : "a"
+>"a" : "a"
+
+invoke({
+>invoke({    kind,    method(a) {        return +a;    }}) : void
+>invoke : (item: ADT) => void
+>{    kind,    method(a) {        return +a;    }} : { kind: "a"; method(a: string): number; }
+
+    kind,
+>kind : "a"
+
+    method(a) {
+>method : (a: string) => number
+>a : string
+
+        return +a;
+>+a : number
+>a : string
+    }
+})
+

tests/cases/compiler/contextuallyTypedByDiscriminableUnion.ts

@@ -23,3 +23,11 @@ invoke({
         return +a;
     }
 });
+
+const kind = "a"
+invoke({
+    kind,
+    method(a) {
+        return +a;
+    }
+})