Better typings for Array.concat() and .flatMap()

Author: jablko

src/harness/fourslash.ts

@@ -4685,6 +4685,7 @@ namespace FourSlashInterface {
             varEntry("JSON"),
             interfaceEntry("ReadonlyArray"),
             interfaceEntry("ConcatArray"),
+            typeEntry("ConcatFlatten"),
             varEntry("Array"),
             interfaceEntry("ArrayConstructor"),
             interfaceEntry("TypedPropertyDescriptor"),

src/lib/es2019.array.d.ts

@@ -11,9 +11,9 @@ interface ReadonlyArray<T> {
      * thisArg is omitted, undefined is used as the this value.
      */
     flatMap<U, This = undefined> (
-        callback: (this: This, value: T, index: number, array: T[]) => U | ReadonlyArray<U>,
+        callback: (this: This, value: T, index: number, array: T[]) => U,
         thisArg?: This
-    ): U[]
+    ): (U extends readonly (infer V)[] ? V : U)[]
 
 
     /**
@@ -125,9 +125,9 @@ interface Array<T> {
      * thisArg is omitted, undefined is used as the this value.
      */
     flatMap<U, This = undefined> (
-        callback: (this: This, value: T, index: number, array: T[]) => U | ReadonlyArray<U>,
+        callback: (this: This, value: T, index: number, array: T[]) => U,
         thisArg?: This
-    ): U[]
+    ): (U extends readonly (infer V)[] ? V : U)[]
 
     /**
      * Returns a new array with all sub-array elements concatenated into it recursively up to the

src/lib/es5.d.ts

@@ -1085,7 +1085,7 @@ interface ReadonlyArray<T> {
      * Combines two or more arrays.
      * @param items Additional items to add to the end of array1.
      */
-    concat(...items: (T | ConcatArray<T>)[]): T[];
+    concat<U extends any[]>(...items: U): (T | ConcatFlatten<U[number]>)[];
     /**
      * Adds all the elements of an array separated by the specified separator string.
      * @param separator A string used to separate one element of an array from the next in the resulting String. If omitted, the array elements are separated with a comma.
@@ -1181,6 +1181,7 @@ interface ConcatArray<T> {
     join(separator?: string): string;
     slice(start?: number, end?: number): T[];
 }
+type ConcatFlatten<T> = T extends ConcatArray<infer U> ? U : T;
 
 interface Array<T> {
     /**
@@ -1213,7 +1214,7 @@ interface Array<T> {
      * Combines two or more arrays.
      * @param items Additional items to add to the end of array1.
      */
-    concat(...items: (T | ConcatArray<T>)[]): T[];
+    concat<U extends any[]>(...items: U): (T | ConcatFlatten<U[number]>)[];
     /**
      * Adds all the elements of an array separated by the specified separator string.
      * @param separator A string used to separate one element of an array from the next in the resulting String. If omitted, the array elements are separated with a comma.

tests/baselines/reference/arrayConcat2.types

@@ -5,17 +5,17 @@ var a: string[] = [];
 
 a.concat("hello", 'world');
 >a.concat("hello", 'world') : string[]
->a.concat : { (...items: ConcatArray<string>[]): string[]; (...items: (string | ConcatArray<string>)[]): string[]; }
+>a.concat : { (...items: ConcatArray<string>[]): string[]; <U extends any[]>(...items: U): (string | ConcatFlatten<U[number]>)[]; }
 >a : string[]
->concat : { (...items: ConcatArray<string>[]): string[]; (...items: (string | ConcatArray<string>)[]): string[]; }
+>concat : { (...items: ConcatArray<string>[]): string[]; <U extends any[]>(...items: U): (string | ConcatFlatten<U[number]>)[]; }
 >"hello" : "hello"
 >'world' : "world"
 
 a.concat('Hello');
 >a.concat('Hello') : string[]
->a.concat : { (...items: ConcatArray<string>[]): string[]; (...items: (string | ConcatArray<string>)[]): string[]; }
+>a.concat : { (...items: ConcatArray<string>[]): string[]; <U extends any[]>(...items: U): (string | ConcatFlatten<U[number]>)[]; }
 >a : string[]
->concat : { (...items: ConcatArray<string>[]): string[]; (...items: (string | ConcatArray<string>)[]): string[]; }
+>concat : { (...items: ConcatArray<string>[]): string[]; <U extends any[]>(...items: U): (string | ConcatFlatten<U[number]>)[]; }
 >'Hello' : "Hello"
 
 var b = new Array<string>();
@@ -25,8 +25,8 @@ var b = new Array<string>();
 
 b.concat('hello');
 >b.concat('hello') : string[]
->b.concat : { (...items: ConcatArray<string>[]): string[]; (...items: (string | ConcatArray<string>)[]): string[]; }
+>b.concat : { (...items: ConcatArray<string>[]): string[]; <U extends any[]>(...items: U): (string | ConcatFlatten<U[number]>)[]; }
 >b : string[]
->concat : { (...items: ConcatArray<string>[]): string[]; (...items: (string | ConcatArray<string>)[]): string[]; }
+>concat : { (...items: ConcatArray<string>[]): string[]; <U extends any[]>(...items: U): (string | ConcatFlatten<U[number]>)[]; }
 >'hello' : "hello"
 

tests/baselines/reference/arrayConcat3.types

@@ -11,9 +11,9 @@ function doStuff<T extends object, T1 extends T>(a: Array<Fn<T>>, b: Array<Fn<T1
 
     b.concat(a);
 >b.concat(a) : Fn<T1>[]
->b.concat : { (...items: ConcatArray<Fn<T1>>[]): Fn<T1>[]; (...items: (Fn<T1> | ConcatArray<Fn<T1>>)[]): Fn<T1>[]; }
+>b.concat : { (...items: ConcatArray<Fn<T1>>[]): Fn<T1>[]; <U extends any[]>(...items: U): (Fn<T1> | ConcatFlatten<U[number]>)[]; }
 >b : Fn<T1>[]
->concat : { (...items: ConcatArray<Fn<T1>>[]): Fn<T1>[]; (...items: (Fn<T1> | ConcatArray<Fn<T1>>)[]): Fn<T1>[]; }
+>concat : { (...items: ConcatArray<Fn<T1>>[]): Fn<T1>[]; <U extends any[]>(...items: U): (Fn<T1> | ConcatFlatten<U[number]>)[]; }
 >a : Fn<T>[]
 }
 

tests/baselines/reference/arrayConcatMap.types

@@ -4,9 +4,9 @@ var x = [].concat([{ a: 1 }], [{ a: 2 }])
 >[].concat([{ a: 1 }], [{ a: 2 }])          .map(b => b.a) : any[]
 >[].concat([{ a: 1 }], [{ a: 2 }])          .map : <U>(callbackfn: (value: any, index: number, array: any[]) => U, thisArg?: any) => U[]
 >[].concat([{ a: 1 }], [{ a: 2 }]) : any[]
->[].concat : { (...items: ConcatArray<any>[]): any[]; (...items: any[]): any[]; }
+>[].concat : { (...items: ConcatArray<any>[]): any[]; <U extends any[]>(...items: U): any[]; }
 >[] : undefined[]
->concat : { (...items: ConcatArray<any>[]): any[]; (...items: any[]): any[]; }
+>concat : { (...items: ConcatArray<any>[]): any[]; <U extends any[]>(...items: U): any[]; }
 >[{ a: 1 }] : { a: number; }[]
 >{ a: 1 } : { a: number; }
 >a : number

tests/baselines/reference/arrayFlatMap.types

@@ -9,17 +9,17 @@ const readonlyArray: ReadonlyArray<number> = [];
 
 array.flatMap((): ReadonlyArray<number> => []); // ok
 >array.flatMap((): ReadonlyArray<number> => []) : number[]
->array.flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U | readonly U[], thisArg?: This) => U[]
+>array.flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U, thisArg?: This) => (U extends readonly (infer V)[] ? V : U)[]
 >array : number[]
->flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U | readonly U[], thisArg?: This) => U[]
+>flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U, thisArg?: This) => (U extends readonly (infer V)[] ? V : U)[]
 >(): ReadonlyArray<number> => [] : () => readonly number[]
 >[] : undefined[]
 
 readonlyArray.flatMap((): ReadonlyArray<number> => []); // ok
 >readonlyArray.flatMap((): ReadonlyArray<number> => []) : number[]
->readonlyArray.flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U | readonly U[], thisArg?: This) => U[]
+>readonlyArray.flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U, thisArg?: This) => (U extends readonly (infer V)[] ? V : U)[]
 >readonlyArray : readonly number[]
->flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U | readonly U[], thisArg?: This) => U[]
+>flatMap : <U, This = undefined>(callback: (this: This, value: number, index: number, array: number[]) => U, thisArg?: This) => (U extends readonly (infer V)[] ? V : U)[]
 >(): ReadonlyArray<number> => [] : () => readonly number[]
 >[] : undefined[]
 

tests/baselines/reference/arrayOfSubtypeIsAssignableToReadonlyArray.errors.txt

@@ -1,7 +1,7 @@
 tests/cases/compiler/arrayOfSubtypeIsAssignableToReadonlyArray.ts(13,1): error TS2322: Type 'A[]' is not assignable to type 'readonly B[]'.
   Property 'b' is missing in type 'A' but required in type 'B'.
 tests/cases/compiler/arrayOfSubtypeIsAssignableToReadonlyArray.ts(18,1): error TS2322: Type 'C<A>' is not assignable to type 'readonly B[]'.
-  The types returned by 'concat(...)' are incompatible between these types.
+  The types returned by 'slice(...)' are incompatible between these types.
     Type 'A[]' is not assignable to type 'B[]'.
       Type 'A' is not assignable to type 'B'.
 
@@ -31,7 +31,7 @@ tests/cases/compiler/arrayOfSubtypeIsAssignableToReadonlyArray.ts(18,1): error T
     rrb = cra; // error: 'A' is not assignable to 'B'
     ~~~
 !!! error TS2322: Type 'C<A>' is not assignable to type 'readonly B[]'.
-!!! error TS2322:   The types returned by 'concat(...)' are incompatible between these types.
+!!! error TS2322:   The types returned by 'slice(...)' are incompatible between these types.
 !!! error TS2322:     Type 'A[]' is not assignable to type 'B[]'.
 !!! error TS2322:       Type 'A' is not assignable to type 'B'.
 

tests/baselines/reference/concatError.types

@@ -15,19 +15,19 @@ fa = fa.concat([0]);
 >fa = fa.concat([0]) : number[]
 >fa : number[]
 >fa.concat([0]) : number[]
->fa.concat : { (...items: ConcatArray<number>[]): number[]; (...items: (number | ConcatArray<number>)[]): number[]; }
+>fa.concat : { (...items: ConcatArray<number>[]): number[]; <U extends any[]>(...items: U): (number | ConcatFlatten<U[number]>)[]; }
 >fa : number[]
->concat : { (...items: ConcatArray<number>[]): number[]; (...items: (number | ConcatArray<number>)[]): number[]; }
+>concat : { (...items: ConcatArray<number>[]): number[]; <U extends any[]>(...items: U): (number | ConcatFlatten<U[number]>)[]; }
 >[0] : number[]
 >0 : 0
 
 fa = fa.concat(0);
 >fa = fa.concat(0) : number[]
 >fa : number[]
 >fa.concat(0) : number[]
->fa.concat : { (...items: ConcatArray<number>[]): number[]; (...items: (number | ConcatArray<number>)[]): number[]; }
+>fa.concat : { (...items: ConcatArray<number>[]): number[]; <U extends any[]>(...items: U): (number | ConcatFlatten<U[number]>)[]; }
 >fa : number[]
->concat : { (...items: ConcatArray<number>[]): number[]; (...items: (number | ConcatArray<number>)[]): number[]; }
+>concat : { (...items: ConcatArray<number>[]): number[]; <U extends any[]>(...items: U): (number | ConcatFlatten<U[number]>)[]; }
 >0 : 0
 
 

tests/baselines/reference/concatTuples.types

@@ -10,9 +10,9 @@ ijs = ijs.concat([[3, 4], [5, 6]]);
 >ijs = ijs.concat([[3, 4], [5, 6]]) : [number, number][]
 >ijs : [number, number][]
 >ijs.concat([[3, 4], [5, 6]]) : [number, number][]
->ijs.concat : { (...items: ConcatArray<[number, number]>[]): [number, number][]; (...items: ([number, number] | ConcatArray<[number, number]>)[]): [number, number][]; }
+>ijs.concat : { (...items: ConcatArray<[number, number]>[]): [number, number][]; <U extends any[]>(...items: U): ([number, number] | ConcatFlatten<U[number]>)[]; }
 >ijs : [number, number][]
->concat : { (...items: ConcatArray<[number, number]>[]): [number, number][]; (...items: ([number, number] | ConcatArray<[number, number]>)[]): [number, number][]; }
+>concat : { (...items: ConcatArray<[number, number]>[]): [number, number][]; <U extends any[]>(...items: U): ([number, number] | ConcatFlatten<U[number]>)[]; }
 >[[3, 4], [5, 6]] : [number, number][]
 >[3, 4] : [number, number]
 >3 : 3