Typescript Generic for cartesian product

38 views Asked by At

Here is a typescript Cartesian product iterator:

export function* productIterator3<T1, T2, T3>(
  ...arr: [T1[], T2[], T3[]]
): Generator<[T1, T2, T3]> {
  const lengths = arr.map((component) => component.length);
  const total = lengths.reduce((prod, len) => prod * len, 1);
  console.log(lengths, total);
  for (let totalIndex = 0; totalIndex < total; totalIndex++) {
    const product = [];
    let reducedIndex = totalIndex;
    for (const iter of arr) {
      product.push(iter[reducedIndex % iter.length]);
      reducedIndex = Math.floor(reducedIndex / iter.length);
    }
    yield product as [T1, T2, T3];
  }
}

type Colors = 'blue' | 'red';
const colors: Colors[] = ['blue', 'red'];
type Sizes = 'small' | 'medium' | 'large';
const sizes: Sizes[] = ['small', 'medium', 'large'];
type Fits = 'turtleneck' | 'tank';
const fits: Fits[] = ['turtleneck', 'tank'];

for (const [color, size, fit] of productIterator3(colors, sizes, fits)) {
  console.log(color, size, fit); // color, size, fit are properly typed
}

The code itself isn't specific to the length of ...arr, but the typescript annotations are. How do I adapt this to be fully generic? This isn't the correct syntax of course, but hopefully it gets the point across:

export function* productIterator<...T>(...arr: [...T[]]): Generator<[...T]> {
1

There are 1 answers

0
Kaia On

The key challenge here is finding a way to go between, for instance:

type TupleOfScalars = [string, number, bool];
type TupleOfArrays = [string[], number[], bool[]];

We can use an index signature using "in" for a union type like so and then index into our type with that K.

type ArraysOfTuple<TupleOfScalars extends ReadonlyArray<unknown>> = {
  [K in keyof TupleOfScalars]: K extends number ? TupleOfScalars[K][] : never;
};

type ScalarsOfTuple<TupleOfArrays extends ReadonlyArray<unknown[]>> = {
  [K in keyof TupleOfArrays]: K extends number ? TupleOfArrays[K][number] : never;
};

type test1 = ArraysOfTuple<[number, string, boolean]>; // [number[], string[], boolean[]]
type test2 = ScalarsOfTuple<[number[], string[], boolean[]]>; // [number, string, boolean]

To break this down:

If we have a union type we can use it as an index for our type.

type Indices = 1 | 2 | 3 | 'someString';
type ThreeTuple = { [K in Indices]: number }; // { 3: number; 1: number; 2: number; someString: number; }

We can use K within that index to further specify the type of that particular attribute. [K in Indices]: K, for instance, would make it so obj[1] has the type 1.

keyof number[] includes 0, 1, 2, etc., but it also includes 'toReversed' and other array functions. So if we just do:

type ArraysOfTuple<TupleOfScalars extends ReadonlyArray<unknown>> = {
  [K in keyof TupleOfScalars]: TupleOfScalars[K][]
};

We'll have [string, number] => [string[], number[]] as desired, but it could mess with the other properties. The ternary operator selects only the numeric indices to modify. never in the else clause drops the item from the ternary. See mapped types.

My final signature:

export function* productIterator<T extends ReadonlyArray<unknown[]>>(
  ...arr: T
): Generator<{ [K in keyof T]: K extends number ? T[K][number] : T[K] }> {

If you didn't have the numeric key distinction, T[K] extends (infer U)[] (as seen in the other answer) is T[K] extends any[] but will capture the type U for use in the ternary.