Fix logic for determining whether to simplify keyof on mapped types

src/compiler/checker.ts

@@ -14250,16 +14250,23 @@ 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 && (!(<ConditionalType>type).root.isDistributive || maybeNonDistributiveNameType((<ConditionalType>type).checkType)) ||
-                type.flags & (TypeFlags.UnionOrIntersection | TypeFlags.TemplateLiteral) && some((<UnionOrIntersectionType | TemplateLiteralType>type).types, maybeNonDistributiveNameType) ||
-                type.flags & (TypeFlags.Index | TypeFlags.StringMapping) && maybeNonDistributiveNameType((<IndexType | StringMappingType>type).type) ||
-                type.flags & TypeFlags.IndexedAccess && maybeNonDistributiveNameType((<IndexedAccessType>type).indexType) ||
-                type.flags & TypeFlags.Substitution && maybeNonDistributiveNameType((<SubstitutionType>type).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 isDistributiveForTypeParameter(getNameTypeFromMappedType(mappedType) || typeVariable);
+            function isDistributiveForTypeParameter(type: Type): boolean {
+                return type.flags & TypeFlags.TypeParameter ? type === typeVariable :
+                    type.flags & TypeFlags.Conditional ? (<ConditionalType>type).root.isDistributive && isDistributiveForTypeParameter((<ConditionalType>type).checkType) :
+                    type.flags & (TypeFlags.UnionOrIntersection | TypeFlags.TemplateLiteral) ? every((<UnionOrIntersectionType | TemplateLiteralType>type).types, isDistributiveForTypeParameter) :
+                    type.flags & TypeFlags.IndexedAccess ? isDistributiveForTypeParameter((<IndexedAccessType>type).indexType) :
+                    type.flags & TypeFlags.Substitution ? isDistributiveForTypeParameter((<SubstitutionType>type).substitute) :
+                    type.flags & TypeFlags.StringMapping ? isDistributiveForTypeParameter((<StringMappingType>type).type) :
+                    type.flags & TypeFlags.Index ? false :
+                    true;
+                }
         }
 
         function getLiteralTypeFromPropertyName(name: PropertyName) {
@@ -14312,7 +14319,7 @@ namespace ts {
             type = getReducedType(type);
             return type.flags & TypeFlags.Union ? getIntersectionType(map((<IntersectionType>type).types, t => getIndexType(t, stringsOnly, noIndexSignatures))) :
                 type.flags & TypeFlags.Intersection ? getUnionType(map((<IntersectionType>type).types, t => getIndexType(t, stringsOnly, noIndexSignatures))) :
-                type.flags & TypeFlags.InstantiableNonPrimitive || isGenericTupleType(type) || isGenericMappedType(type) && maybeNonDistributiveNameType(getNameTypeFromMappedType(type)) ? getIndexTypeForGenericType(<InstantiableType | UnionOrIntersectionType>type, stringsOnly) :
+                type.flags & TypeFlags.InstantiableNonPrimitive || isGenericTupleType(type) || isGenericMappedType(type) && !hasDistributiveNameType(type) ? getIndexTypeForGenericType(<InstantiableType | UnionOrIntersectionType>type, stringsOnly) :
                 getObjectFlags(type) & ObjectFlags.Mapped ? getIndexTypeForMappedType(<MappedType>type, noIndexSignatures) :
                 type === wildcardType ? wildcardType :
                 type.flags & TypeFlags.Unknown ? neverType :

tests/baselines/reference/mappedTypeAsClauses.errors.txt

@@ -0,0 +1,137 @@
+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"'.
+

tests/baselines/reference/mappedTypeAsClauses.js

@@ -114,6 +114,21 @@ 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"
 
 
 //// [mappedTypeAsClauses.js]
@@ -126,6 +141,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 +260,11 @@ 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";

tests/baselines/reference/mappedTypeAsClauses.symbols

@@ -369,3 +369,41 @@ type Res3 = keyof GetObjWithIf<Schema, Schema['root']['task']>;  // "Task"
 >Schema : Symbol(Schema, Decl(mappedTypeAsClauses.ts, 102, 2))
 >Schema : Symbol(Schema, Decl(mappedTypeAsClauses.ts, 102, 2))
 
+// Repro from #44019
+
+type KeysExtendedBy<T, U> = keyof { [K in keyof T as U extends T[K] ? K : never] : T[K] };
+>KeysExtendedBy : Symbol(KeysExtendedBy, Decl(mappedTypeAsClauses.ts, 114, 63))
+>T : Symbol(T, Decl(mappedTypeAsClauses.ts, 118, 20))
+>U : Symbol(U, Decl(mappedTypeAsClauses.ts, 118, 22))
+>K : Symbol(K, Decl(mappedTypeAsClauses.ts, 118, 37))
+>T : Symbol(T, Decl(mappedTypeAsClauses.ts, 118, 20))
+>U : Symbol(U, Decl(mappedTypeAsClauses.ts, 118, 22))
+>T : Symbol(T, Decl(mappedTypeAsClauses.ts, 118, 20))
+>K : Symbol(K, Decl(mappedTypeAsClauses.ts, 118, 37))
+>K : Symbol(K, Decl(mappedTypeAsClauses.ts, 118, 37))
+>T : Symbol(T, Decl(mappedTypeAsClauses.ts, 118, 20))
+>K : Symbol(K, Decl(mappedTypeAsClauses.ts, 118, 37))
+
+interface M {
+>M : Symbol(M, Decl(mappedTypeAsClauses.ts, 118, 90))
+
+    a: boolean;
+>a : Symbol(M.a, Decl(mappedTypeAsClauses.ts, 120, 13))
+
+    b: number;
+>b : Symbol(M.b, Decl(mappedTypeAsClauses.ts, 121, 15))
+}
+
+function f(x: KeysExtendedBy<M, number>) {
+>f : Symbol(f, Decl(mappedTypeAsClauses.ts, 123, 1))
+>x : Symbol(x, Decl(mappedTypeAsClauses.ts, 125, 11))
+>KeysExtendedBy : Symbol(KeysExtendedBy, Decl(mappedTypeAsClauses.ts, 114, 63))
+>M : Symbol(M, Decl(mappedTypeAsClauses.ts, 118, 90))
+
+    return x;
+>x : Symbol(x, Decl(mappedTypeAsClauses.ts, 125, 11))
+}
+
+f("a");  // Error, should allow only "b"
+>f : Symbol(f, Decl(mappedTypeAsClauses.ts, 123, 1))
+

tests/baselines/reference/mappedTypeAsClauses.types

@@ -225,3 +225,29 @@ type Res2 = GetKeyWithIf<Schema, Schema['root']['task']>;  // "Task"
 type Res3 = keyof GetObjWithIf<Schema, Schema['root']['task']>;  // "Task"
 >Res3 : "Task"
 
+// Repro from #44019
+
+type KeysExtendedBy<T, U> = keyof { [K in keyof T as U extends T[K] ? K : never] : T[K] };
+>KeysExtendedBy : keyof { [K in keyof T as U extends T[K] ? K : never]: T[K]; }
+
+interface M {
+    a: boolean;
+>a : boolean
+
+    b: number;
+>b : number
+}
+
+function f(x: KeysExtendedBy<M, number>) {
+>f : (x: KeysExtendedBy<M, number>) => "b"
+>x : "b"
+
+    return x;
+>x : "b"
+}
+
+f("a");  // Error, should allow only "b"
+>f("a") : "b"
+>f : (x: "b") => "b"
+>"a" : "a"
+

tests/cases/conformance/types/mapped/mappedTypeAsClauses.ts

@@ -116,3 +116,18 @@ 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"