Fix logic for determining whether to simplify keyof on mapped types (#…

…44042)

* Fix logic for determining whether to simplify keyof on mapped types

* Add regression test

* Improve hasDistributiveNameType check

* Add more tests

* Address code review feedback

* Add more tests

src/compiler/checker.ts

@@ -14333,16 +14333,22 @@ namespace ts {
         }
 
         // Ordinarily we reduce a keyof M, where M is a mapped type { [P in K as N<P>]: X }, to simply N<K>. This however presumes
-        // that N distributes over union types, i.e. that N<A | B | C> is equivalent to N<A> | N<B> | N<C>. That presumption may not
-        // be true when N is a non-distributive conditional type or an instantiable type with a non-distributive conditional type as
-        // a constituent. In those cases, we cannot reduce keyof M and need to preserve it as is.
-        function maybeNonDistributiveNameType(type: Type | undefined): boolean {
-            return !!(type && (
-                type.flags & TypeFlags.Conditional && (!(type as ConditionalType).root.isDistributive || maybeNonDistributiveNameType((type as ConditionalType).checkType)) ||
-                type.flags & (TypeFlags.UnionOrIntersection | TypeFlags.TemplateLiteral) && some((type as UnionOrIntersectionType | TemplateLiteralType).types, maybeNonDistributiveNameType) ||
-                type.flags & (TypeFlags.Index | TypeFlags.StringMapping) && maybeNonDistributiveNameType((type as IndexType | StringMappingType).type) ||
-                type.flags & TypeFlags.IndexedAccess && maybeNonDistributiveNameType((type as IndexedAccessType).indexType) ||
-                type.flags & TypeFlags.Substitution && maybeNonDistributiveNameType((type as SubstitutionType).substitute)));
+        // that N distributes over union types, i.e. that N<A | B | C> is equivalent to N<A> | N<B> | N<C>. Specifically, we only
+        // want to perform the reduction when the name type of a mapped type is distributive with respect to the type variable
+        // introduced by the 'in' clause of the mapped type. Note that non-generic types are considered to be distributive because
+        // they're the same type regardless of what's being distributed over.
+        function hasDistributiveNameType(mappedType: MappedType) {
+            const typeVariable = getTypeParameterFromMappedType(mappedType);
+            return isDistributive(getNameTypeFromMappedType(mappedType) || typeVariable);
+            function isDistributive(type: Type): boolean {
+                return type.flags & (TypeFlags.AnyOrUnknown | TypeFlags.Primitive | TypeFlags.Never | TypeFlags.TypeParameter | TypeFlags.Object | TypeFlags.NonPrimitive) ? true :
+                    type.flags & TypeFlags.Conditional ? (type as ConditionalType).root.isDistributive && (type as ConditionalType).checkType === typeVariable :
+                    type.flags & (TypeFlags.UnionOrIntersection | TypeFlags.TemplateLiteral) ? every((type as UnionOrIntersectionType | TemplateLiteralType).types, isDistributive) :
+                    type.flags & TypeFlags.IndexedAccess ? isDistributive((type as IndexedAccessType).objectType) && isDistributive((type as IndexedAccessType).indexType) :
+                    type.flags & TypeFlags.Substitution ? isDistributive((type as SubstitutionType).substitute) :
+                    type.flags & TypeFlags.StringMapping ? isDistributive((type as StringMappingType).type) :
+                    false;
+            }
         }
 
         function getLiteralTypeFromPropertyName(name: PropertyName) {
@@ -14393,9 +14399,9 @@ namespace ts {
         function getIndexType(type: Type, stringsOnly = keyofStringsOnly, noIndexSignatures?: boolean): Type {
             const includeOrigin = stringsOnly === keyofStringsOnly && !noIndexSignatures;
             type = getReducedType(type);
-            return type.flags & TypeFlags.Union ? getIntersectionType(map((type as IntersectionType).types, t => getIndexType(t, stringsOnly, noIndexSignatures))) :
+            return type.flags & TypeFlags.Union ? getIntersectionType(map((type as UnionType).types, t => getIndexType(t, stringsOnly, noIndexSignatures))) :
                 type.flags & TypeFlags.Intersection ? getUnionType(map((type as IntersectionType).types, t => getIndexType(t, stringsOnly, noIndexSignatures))) :
-                type.flags & TypeFlags.InstantiableNonPrimitive || isGenericTupleType(type) || isGenericMappedType(type) && maybeNonDistributiveNameType(getNameTypeFromMappedType(type)) ? getIndexTypeForGenericType(type as InstantiableType | UnionOrIntersectionType, stringsOnly) :
+                type.flags & TypeFlags.InstantiableNonPrimitive || isGenericTupleType(type) || isGenericMappedType(type) && !hasDistributiveNameType(type) ? getIndexTypeForGenericType(type as InstantiableType | UnionOrIntersectionType, stringsOnly) :
                 getObjectFlags(type) & ObjectFlags.Mapped ? getIndexTypeForMappedType(type as MappedType, noIndexSignatures) :
                 type === wildcardType ? wildcardType :
                 type.flags & TypeFlags.Unknown ? neverType :

tests/baselines/reference/mappedTypeAsClauses.errors.txt

@@ -0,0 +1,158 @@
+tests/cases/conformance/types/mapped/mappedTypeAsClauses.ts(130,3): error TS2345: Argument of type '"a"' is not assignable to parameter of type '"b"'.
+
+
+==== tests/cases/conformance/types/mapped/mappedTypeAsClauses.ts (1 errors) ====
+    // Mapped type 'as N' clauses
+
+    type Getters<T> = { [P in keyof T & string as `get${Capitalize<P>}`]: () => T[P] };
+    type TG1 = Getters<{ foo: string, bar: number, baz: { z: boolean } }>;
+
+    // Mapped type with 'as N' clause has no constraint on 'in T' clause
+
+    type PropDef<K extends keyof any, T> = { name: K, type: T };
+
+    type TypeFromDefs<T extends PropDef<keyof any, any>> = { [P in T as P['name']]: P['type'] };
+
+    type TP1 = TypeFromDefs<{ name: 'a', type: string } | { name: 'b', type: number } | { name: 'a', type: boolean }>;
+
+    // No array or tuple type mapping when 'as N' clause present
+
+    type TA1 = Getters<string[]>;
+    type TA2 = Getters<[number, boolean]>;
+
+    // Filtering using 'as N' clause
+
+    type Methods<T> = { [P in keyof T as T[P] extends Function ? P : never]: T[P] };
+    type TM1 = Methods<{ foo(): number, bar(x: string): boolean, baz: string | number }>;
+
+    // Mapping to multiple names using 'as N' clause
+
+    type DoubleProp<T> = { [P in keyof T & string as `${P}1` | `${P}2`]: T[P] }
+    type TD1 = DoubleProp<{ a: string, b: number }>;  // { a1: string, a2: string, b1: number, b2: number }
+    type TD2 = keyof TD1;  // 'a1' | 'a2' | 'b1' | 'b2'
+    type TD3<U> = keyof DoubleProp<U>;  // `${keyof U & string}1` | `${keyof U & string}2`
+
+    // Repro from #40619
+
+    type Lazyify<T> = {
+        [K in keyof T as `get${Capitalize<K & string>}`]: () => T[K]
+    };
+
+    interface Person {
+        readonly name: string;
+        age: number;
+        location?: string;
+    }
+
+    type LazyPerson = Lazyify<Person>;
+
+    // Repro from #40833
+
+    type Example = {foo: string, bar: number};
+
+    type PickByValueType<T, U> = {
+      [K in keyof T as T[K] extends U ? K : never]: T[K]
+    };
+
+    type T1 = PickByValueType<Example, string>;
+    const e1: T1 = {
+        foo: "hello"
+    };
+    type T2 = keyof T1;
+    const e2: T2 = "foo";
+
+    // Repro from #41133
+
+    interface Car {
+        name: string;
+        seats: number;
+        engine: Engine;
+        wheels: Wheel[];
+    }
+
+    interface Engine {
+        manufacturer: string;
+        horsepower: number;
+    }
+
+    interface Wheel {
+        type: "summer" | "winter";
+        radius: number;
+    }
+
+    type Primitive = string | number | boolean;
+    type OnlyPrimitives<T> = { [K in keyof T as T[K] extends Primitive ? K : never]: T[K] };
+
+    let primitiveCar: OnlyPrimitives<Car>;  // { name: string; seats: number; }
+    let keys: keyof OnlyPrimitives<Car>;  //  "name" | "seats"
+
+    type KeysOfPrimitives<T> = keyof OnlyPrimitives<T>;
+
+    let carKeys: KeysOfPrimitives<Car>;  // "name" | "seats"
+
+    // Repro from #41453
+
+    type Equal<A, B> = (<T>() => T extends A ? 1 : 2) extends (<T>() => T extends B ? 1 : 2) ? true : false;
+
+    type If<Cond extends boolean, Then, Else> = Cond extends true ? Then : Else;
+
+    type GetKey<S, V> = keyof { [TP in keyof S as Equal<S[TP], V> extends true ? TP : never]: any };
+
+    type GetKeyWithIf<S, V> = keyof { [TP in keyof S as If<Equal<S[TP], V>, TP, never>]: any };
+
+    type GetObjWithIf<S, V> = { [TP in keyof S as If<Equal<S[TP], V>, TP, never>]: any };
+
+    type Task = {
+      isDone: boolean;
+    };
+
+    type Schema = {
+      root: {
+        title: string;
+        task: Task;
+      }
+      Task: Task;
+    };
+
+    type Res1 = GetKey<Schema, Schema['root']['task']>;  // "Task"
+    type Res2 = GetKeyWithIf<Schema, Schema['root']['task']>;  // "Task"
+    type Res3 = keyof GetObjWithIf<Schema, Schema['root']['task']>;  // "Task"
+
+    // Repro from #44019
+
+    type KeysExtendedBy<T, U> = keyof { [K in keyof T as U extends T[K] ? K : never] : T[K] };
+
+    interface M {
+        a: boolean;
+        b: number;
+    }
+
+    function f(x: KeysExtendedBy<M, number>) {
+        return x;
+    }
+
+    f("a");  // Error, should allow only "b"
+      ~~~
+!!! error TS2345: Argument of type '"a"' is not assignable to parameter of type '"b"'.
+
+    type NameMap = { 'a': 'x', 'b': 'y', 'c': 'z' };
+
+    // Distributive, will be simplified
+
+    type TS0<T> = keyof { [P in keyof T as keyof Record<P, number>]: string };
+    type TS1<T> = keyof { [P in keyof T as Extract<P, 'a' | 'b' | 'c'>]: string };
+    type TS2<T> = keyof { [P in keyof T as P & ('a' | 'b' | 'c')]: string };
+    type TS3<T> = keyof { [P in keyof T as Exclude<P, 'a' | 'b' | 'c'>]: string };
+    type TS4<T> = keyof { [P in keyof T as NameMap[P & keyof NameMap]]: string };
+    type TS5<T> = keyof { [P in keyof T & keyof NameMap as NameMap[P]]: string };
+    type TS6<T, U, V> = keyof { [ K in keyof T as V & (K extends U ? K : never)]: string };
+
+    // Non-distributive, won't be simplified
+
+    type TN0<T> = keyof { [P in keyof T as T[P] extends number ? P : never]: string };
+    type TN1<T> = keyof { [P in keyof T as number extends T[P] ? P : never]: string };
+    type TN2<T> = keyof { [P in keyof T as 'a' extends P ? 'x' : 'y']: string };
+    type TN3<T> = keyof { [P in keyof T as Exclude<Exclude<Exclude<P, 'c'>, 'b'>, 'a'>]: string };
+    type TN4<T, U> = keyof { [K in keyof T as (K extends U ? T[K] : never) extends T[K] ? K : never]: string };
+    type TN5<T, U> = keyof { [K in keyof T as keyof { [P in K as T[P] extends U ? K : never]: true }]: string };
+

tests/baselines/reference/mappedTypeAsClauses.js

@@ -114,6 +114,42 @@ type Schema = {
 type Res1 = GetKey<Schema, Schema['root']['task']>;  // "Task"
 type Res2 = GetKeyWithIf<Schema, Schema['root']['task']>;  // "Task"
 type Res3 = keyof GetObjWithIf<Schema, Schema['root']['task']>;  // "Task"
+
+// Repro from #44019
+
+type KeysExtendedBy<T, U> = keyof { [K in keyof T as U extends T[K] ? K : never] : T[K] };
+
+interface M {
+    a: boolean;
+    b: number;
+}
+
+function f(x: KeysExtendedBy<M, number>) {
+    return x;
+}
+
+f("a");  // Error, should allow only "b"
+
+type NameMap = { 'a': 'x', 'b': 'y', 'c': 'z' };
+
+// Distributive, will be simplified
+
+type TS0<T> = keyof { [P in keyof T as keyof Record<P, number>]: string };
+type TS1<T> = keyof { [P in keyof T as Extract<P, 'a' | 'b' | 'c'>]: string };
+type TS2<T> = keyof { [P in keyof T as P & ('a' | 'b' | 'c')]: string };
+type TS3<T> = keyof { [P in keyof T as Exclude<P, 'a' | 'b' | 'c'>]: string };
+type TS4<T> = keyof { [P in keyof T as NameMap[P & keyof NameMap]]: string };
+type TS5<T> = keyof { [P in keyof T & keyof NameMap as NameMap[P]]: string };
+type TS6<T, U, V> = keyof { [ K in keyof T as V & (K extends U ? K : never)]: string };
+
+// Non-distributive, won't be simplified
+
+type TN0<T> = keyof { [P in keyof T as T[P] extends number ? P : never]: string };
+type TN1<T> = keyof { [P in keyof T as number extends T[P] ? P : never]: string };
+type TN2<T> = keyof { [P in keyof T as 'a' extends P ? 'x' : 'y']: string };
+type TN3<T> = keyof { [P in keyof T as Exclude<Exclude<Exclude<P, 'c'>, 'b'>, 'a'>]: string };
+type TN4<T, U> = keyof { [K in keyof T as (K extends U ? T[K] : never) extends T[K] ? K : never]: string };
+type TN5<T, U> = keyof { [K in keyof T as keyof { [P in K as T[P] extends U ? K : never]: true }]: string };
 
 
 //// [mappedTypeAsClauses.js]
@@ -126,6 +162,10 @@ var e2 = "foo";
 var primitiveCar; // { name: string; seats: number; }
 var keys; //  "name" | "seats"
 var carKeys; // "name" | "seats"
+function f(x) {
+    return x;
+}
+f("a"); // Error, should allow only "b"
 
 
 //// [mappedTypeAsClauses.d.ts]
@@ -241,3 +281,57 @@ declare type Schema = {
 declare type Res1 = GetKey<Schema, Schema['root']['task']>;
 declare type Res2 = GetKeyWithIf<Schema, Schema['root']['task']>;
 declare type Res3 = keyof GetObjWithIf<Schema, Schema['root']['task']>;
+declare type KeysExtendedBy<T, U> = keyof {
+    [K in keyof T as U extends T[K] ? K : never]: T[K];
+};
+interface M {
+    a: boolean;
+    b: number;
+}
+declare function f(x: KeysExtendedBy<M, number>): "b";
+declare type NameMap = {
+    'a': 'x';
+    'b': 'y';
+    'c': 'z';
+};
+declare type TS0<T> = keyof {
+    [P in keyof T as keyof Record<P, number>]: string;
+};
+declare type TS1<T> = keyof {
+    [P in keyof T as Extract<P, 'a' | 'b' | 'c'>]: string;
+};
+declare type TS2<T> = keyof {
+    [P in keyof T as P & ('a' | 'b' | 'c')]: string;
+};
+declare type TS3<T> = keyof {
+    [P in keyof T as Exclude<P, 'a' | 'b' | 'c'>]: string;
+};
+declare type TS4<T> = keyof {
+    [P in keyof T as NameMap[P & keyof NameMap]]: string;
+};
+declare type TS5<T> = keyof {
+    [P in keyof T & keyof NameMap as NameMap[P]]: string;
+};
+declare type TS6<T, U, V> = keyof {
+    [K in keyof T as V & (K extends U ? K : never)]: string;
+};
+declare type TN0<T> = keyof {
+    [P in keyof T as T[P] extends number ? P : never]: string;
+};
+declare type TN1<T> = keyof {
+    [P in keyof T as number extends T[P] ? P : never]: string;
+};
+declare type TN2<T> = keyof {
+    [P in keyof T as 'a' extends P ? 'x' : 'y']: string;
+};
+declare type TN3<T> = keyof {
+    [P in keyof T as Exclude<Exclude<Exclude<P, 'c'>, 'b'>, 'a'>]: string;
+};
+declare type TN4<T, U> = keyof {
+    [K in keyof T as (K extends U ? T[K] : never) extends T[K] ? K : never]: string;
+};
+declare type TN5<T, U> = keyof {
+    [K in keyof T as keyof {
+        [P in K as T[P] extends U ? K : never]: true;
+    }]: string;
+};