Support recursive ReadonlyArray

Author: blakeembrey

src/index.spec.ts

@@ -1,28 +1,59 @@
 import { expectType } from "ts-expect";
 import { flatten } from "./index";
 
-describe("flatten", function() {
-  it("should flatten an array", function() {
+describe("flatten", () => {
+  it("should flatten an array", () => {
     const result = flatten([1, [2, [3, [4, [5]]], 6, [[7], 8], 9], 10]);
 
     expectType<number[]>(result);
 
     expect(result).toStrictEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
   });
 
-  it("should work with array-like", function() {
+  it("should work with array-like", () => {
     const result = flatten("test");
 
     expectType<string[]>(result);
 
     expect(result).toStrictEqual(["t", "e", "s", "t"]);
   });
 
-  it("should work with arguments", function() {
-    const result = flatten(arguments);
+  it("should work with readonly array", () => {
+    const input = [1, [2, [3, [4]]]] as const;
+    const result = flatten(input);
+
+    expectType<(1 | 2 | 3 | 4)[]>(result);
+
+    expect(result).toStrictEqual([1, 2, 3, 4]);
+  });
+
+  it("should work with arguments", () => {
+    const input = (function() {
+      return arguments;
+    })();
+    const result = flatten(input);
 
     expectType<any[]>(result);
 
     expect(result).toStrictEqual([]);
   });
+
+  it("should work with mixed types", () => {
+    const fn = (x: string) => x;
+    const input = [1, ["test", [fn, [true]]]];
+    const result = flatten(input);
+
+    expectType<(number | string | boolean | typeof fn)[]>(result);
+
+    expect(result).toStrictEqual([1, "test", fn, true]);
+  });
+
+  it("should work with tuples", () => {
+    const input: [number, [number, number], [number]] = [1, [1, 2], [3]];
+    const result = flatten(input);
+
+    expectType<number[]>(result);
+
+    expect(result).toStrictEqual([1, 1, 2, 3]);
+  });
 });

src/index.ts

@@ -1,25 +1,40 @@
 /**
- * Recursive value or array type.
+ * Pick the value from an array.
  */
-export type ValueOrArray<T> = T | Array<ValueOrArray<T>>;
+export type PickValue<T> = T extends ReadonlyArray<any>
+  ? {
+      [K in Extract<keyof T, number>]: PickValue<T[K]>;
+    }[number]
+  : T;
+
+/**
+ * Flatten an `ArrayLike` object in TypeScript.
+ */
+export type FlatArray<T extends ArrayLike<any>> = Array<PickValue<T[number]>>;
 
 /**
  * Flatten an array indefinitely.
  */
-export function flatten<T>(array: ArrayLike<ValueOrArray<T>>): T[] {
-  return $flatten(array, []);
+export function flatten<T extends ArrayLike<any>>(array: T): FlatArray<T> {
+  const result: FlatArray<T> = [];
+  $flatten<T>(array, result);
+  return result;
 }
 
-function $flatten<T>(array: ArrayLike<ValueOrArray<T>>, result: T[]): T[] {
+/**
+ * Internal flatten function recursively passes `result`.
+ */
+function $flatten<T extends ArrayLike<any>>(
+  array: T,
+  result: FlatArray<T>
+): void {
   for (let i = 0; i < array.length; i++) {
     const value = array[i];
 
     if (Array.isArray(value)) {
-      $flatten(value, result);
+      $flatten(value as any, result);
     } else {
       result.push(value);
     }
   }
-
-  return result;
 }