Search code examples
typescript

Infer size of smallest tuple if all arrays are tuple, fallback to array otherwise


I'm trying to define a generic zip function:

type ZippedElement<T extends unknown[][]> = {
  [K in keyof T]: T[K] extends (infer V)[] ? V : never;
};

export function zip<T extends unknown[][]>(...args: T) {
  const minLength = Math.min(...args.map((arr) => arr.length));

  return Array.from({ length: minLength }, (_, i) =>
    args.map((arr) => arr[i]),
  ) as ZippedElement<T>[];
}

It seems to work with dynamic arrays:

const a: number[] = ...;
const b: string[] = ...;
zip(a, b) // ==> [number, string][]

But I want my type checking to also infer tuple when possible, so I want this to work too:

const a: [number, number] = [1, 2];
const b: [string, string] = ["a", "b"];
zip(a, b) // ==> [[number, string], [number, string]]

I created myself an utility for creating a tuple of a particular size:

type StaticArray<L extends number, T, R extends any[] = T[]> =
  R extends { length: L }
    ? R
    : StaticArray<L, T, [...R, T]>;

All that would be missing is a type that takes T extends unknown[][] and returns a tuple of the minimum size of all tuple if all parameters are tuple, and return an array if there's one array in the inputs:

type ZippedTupleOrArray<T extends unknown[][]> =
  T extends /* magic infer minimum tuple size N */
    ? StaticArray<N, ZippedElement<T>>
    : Array<ZippedElement<T>>;

Is it actually possible to implement? If so, how?


Solution

  • There are multiple ways to write MinimumLength<T> such that, when T is a fixed-length tuple of fixed-length tuples, it evaluates to the length property of the shortest such tuple. One such way is presented in this answer.


    <CAVEAT>

    Of course, such type juggling tends to have weird edge case behavior, so it's important to test any potential solution against a wide range of use cases to be sure it meets your needs.

    The case you seem to care about is detecting if any of the members are just array types and not tuples.

    Edge cases I'm not concerning myself with in this answer:

    The point is that it's important to test. If you hit an edge case behavior that's unacceptable, be prepared to significantly refactor things.

    </CAVEAT>


    If we have a MimimumLength<T> that works for fixed-length tuples, you can write ZippedTupleOrArray to handle non-tuple types by just checking whether the length property of the union of the elements is number. A fixed-length tuple like [0, 0] has a numeric literal type for length, like 2. A regular array type like string[], on the other hand, has number. And the union of any numeric literal type with number is just number. So if any of the elements of T aren't tuples, then T[number]['length'] will be number. Otherwise it will be some proper subtype of number:

    type ZippedTupleOrArray<T extends unknown[][]> =
        number extends T[number]['length'] ?
        Array<ZippedElement<T>> : StaticArray<MinimumLength<T>, ZippedElement<T>>
    

    Now we can implement MinimumLength<T> and assume all tuples inside T:

    type MinimumLength<T extends any[][], N extends number = number> =
        T extends [infer F extends any[], ...infer R extends any[][]] ?
        MinimumLength<R, `${N}` extends keyof F ? N : F['length']> :
        N
    

    The way this works is to recognize that a tuple has numeric-like string literal keys corresponding to the indices. So keyof ["", "", ""] will include "0", "1", and "2". And if you take the length and stringify it with a template literal type, you'll get a numeric-like string literal type one greater than all of those: so `${["", "", ""]['length']}` is "3".

    So we can walk through T and accumulate the minimum length N (using a tail-recursive conditional type). For each element F of T, we check to see whether the stringified N is a key of F. If it is, then F is longer than N so we keep N as-is. If not, then F is shorter than N, so we toss N and use F['length'] instead.

    We start with an N of number so that MinimumLength<[]> is just number. As soon as we check anything like MinimumLength<[["","",""]]>, then number is first replaced with the length of the first element, because `${number}` is not explicitly a key of a fixed-length tuple.

    And after we walk through all of T and end up with MinimumLength<[], N>, we return the accumulated N.


    Let's test it:

    type Test0 = MinimumLength<[[1], [2, 2], [], [3, 3, 3], [2, 2]]>;
    //   ^? type Test0 = 0
    
    type Test1 = MinimumLength<[[1], [2, 2], [3, 3, 3]]>;
    //   ^? type Test1 = 1
    

    Looks good. And now let's test zip()

    declare function zip<T extends unknown[][]>(...args: T): ZippedTupleOrArray<T>;
    
    declare const a: number[];
    declare const b: string[];
    zip(a, b) // [number, string][]
    
    const c: [number, number] = [1, 2];
    const d: [string, string] = ["a", "b"];
    zip(c, d) // [[number, string], [number, string]]
    
    zip(a, b, c, d); // [number, string, number, string][]
    
    zip(c, d, [null, null, null, null] as const); 
    // [[number, string, null], [number, string, null]]
    

    Also looks good.

    Playground link to code