TS 类型体操笔记 - 869 DistributeUnions

109 阅读6分钟

引言

类型体操 是关于 TS 类型编程的一系列挑战(编程题目)。通过这些挑战,我们可以加深对 TS 类型系统的理解,举一反三,在实际工作中解决类似问题。

本文是我对 869 DistributeUnions 的解题笔记。

本文的读者是正在做 TS 类型体操,对上述题目感兴趣,希望获得更多解读的同学。读者需要具备一定的 TS 基础知识,这部分推荐阅读另一篇优秀博文 Typescript 类型编程,从入门到念头通达😇

题目和答案

先直接贴出题目与答案,呈现问题全貌。

题目

Implement a type Distribute Unions, that turns a type of data structure containing union types into a union of all possible types of permitted data structures that don't contain any union. The data structure can be any combination of objects and tuples on any level of nesting.

For example:

type T1 = DistributeUnions<[1 | 2, 'a' | 'b']>
// =>   [1, 'a'] | [2, 'a'] | [1, 'b'] | [2, 'b']

type T2 = DistributeUnions<{ type: 'a', value: number | string } | { type: 'b', value: boolean }>
//  =>  | { type 'a', value: number }
//      | { type 'a', value: string }
//      | { type 'b', value: boolean }

type T3 = DistributeUnions<[{ value: 'a' | 'b' },  { x: { y: 2 | 3  } }] | 17>
//  =>  | [{ value: 'a' },  { x: { y: 2  } }]
//      | [{ value: 'a' },  { x: { y: 3  } }]
//      | [{ value: 'b' },  { x: { y: 2  } }]
//      | [{ value: 'b' },  { x: { y: 3  } }]
//      | 17

For context, this type can be very useful if you want to exclude a case on deep data structures:

type ExcludeDeep<A, B> = Exclude<DistributeUnions<A>, B>

type T0 = ExcludeDeep<[{ value: 'a' | 'b' },  { x: { y: 2 | 3  } }] | 17, [{ value: 'a' },  any]>
//  =>  | [{ value: 'b' },  { x: { y: 2  } }]
//      | [{ value: 'b' },  { x: { y: 3  } }]
//      | 17

答案

Issue 23086

版本 1:简单有效

type DistributeTuple<T extends unknown[]> = 
  T extends [infer Head, ...infer Rest]
  ? DistributeUnions<Head> extends infer Value
    ? Value extends unknown
      ? [Value, ...DistributeTuple<Rest>]
      : never
    : never
  : T

type DistributeRecord<T extends Record<PropertyKey, unknown>, _Key extends keyof T = keyof T> = 
  T extends Record<PropertyKey, never> ? {} : 
  _Key extends unknown
  ? DistributeUnions<T[_Key]> extends infer Value
    ? Value extends unknown
      ? DistributeRecord<Omit<T, _Key>> extends infer Rest
        ? Rest extends unknown
          ? {[P in keyof T]: P extends keyof Rest ? Rest[P] : Value}
          : never
        : never
      : never
    : never
  : never

type DistributeUnions<T> = 
  T extends unknown[]
  ? DistributeTuple<T>
  : T extends Record<PropertyKey, unknown>
    ? DistributeRecord<T>
    : T

版本 2:算法更高效

type UnionToIntersection<T> = (T extends any ? (x: T) => void : never) extends (x: infer U) => void ?  U : never
type PopUnion<T> = UnionToIntersection<T extends any ? (x: T) => void : never> extends (x: infer U) => void ? U : never
type OneAndRest<T> = [PopUnion<T>, Exclude<T, PopUnion<T>>]

type DistributeTuple<T extends unknown[]> = 
  T extends [infer Head, ...infer Rest]
  ? DistributeUnions<Head> extends infer Value
    ? Value extends unknown
      ? [Value, ...DistributeTuple<Rest>]
      : never
    : never
  : T

type DistributeRecord<T extends Record<PropertyKey, unknown>> = 
  T extends Record<PropertyKey, never> ? {} : 
  OneAndRest<keyof T> extends [infer Key extends keyof T, infer RestKeys extends keyof T]
  ? DistributeUnions<T[Key]> extends infer Value
    ? Value extends unknown
      ? DistributeRecord<Pick<T, RestKeys>> extends infer RestRecord
        ? RestRecord extends unknown
          ? {[P in keyof T]: P extends keyof RestRecord ? RestRecord[P] : Value}
          : never
        : never
      : never
    : never
  : never

type DistributeUnions<T> = 
  T extends unknown[]
  ? DistributeTuple<T>
  : T extends Record<PropertyKey, unknown>
    ? DistributeRecord<T>
    : T

解题思路

首先来理解一下问题。这个题目其实挺像某种 flatten 类型的问题,我们需要对数据结构进行深度遍历,然后将内部的信息提取到顶层并摊平。因此,我们可能需要用上 递归深度遍历 等手段。

理解问题之后,我们首先来设计一下算法。

对于元组数据结构

对于元组,我们顺着它的 index 来进行 reduce。我们可以把一个元组解析为 head 和 rest 两部分,然后将它们各自映射为不同的结果,最后将这两个结果以某种方式合并起来。

image.png

注意,上述结果 A 和结果 B 都是 union,最终的结果也是个 union,并且是 A 和 B 的笛卡尔积。

image.png

在编程中,如果没有其他语法糖工具帮忙的话,可以使用两层迭代就能实现对两个集合做叉乘:

for (const x in [a1, a2]) {
  for (const y in [b1, b2]) {
    // ...
  }
}

在 TS 中,结构是一样的,只是限于语言特性,看起来略有不同:

type DistributeTuple<T extends unknown[]> = 
  T extends [infer Head, ...infer Rest]
  ? DistributeUnions<Head> extends infer Value
    ? Value extends unknown
      ? [Value, ...DistributeTuple<Rest>]
      : never
    : never
  : T

T extends [infer Head, ...infer Rest] 将元组解析为了两个部分: head 和 rest。

代码中看起来比较费解的部分大概是下面这两行:

image

这其实是利用了 TS 中的 分布式条件类型, 你可以把上述代码解读为:

  1. 将 DistributeUnions<Head> 保存为变量 Value,它是一个 union (或者说,数学中的集合)
  2. 然后,对 Value 这个 union 进行迭代 (就像编程中的 for 循环一样), 后续的代码中, Value 不再指原来的 union,而是当前正在被处理的其中一个具体的类型

好了,现在我们有了第一层迭代结构。那说好的第二层呢?放心,我们这里有语法糖呢。在 TS 中,当我们在一个元组中展开一堆元组的 union 时,其结果是一个元组的联合,它们都具有该表达式的相同结构,但每一个元组填充的都是上述被展开的元组 union 们的笛卡尔积中的其中一个成员。读起来有点拗口吼,没事,看下面这张图:

image

好了,理解之后我们再看回来:

image

DistributeTuple<Rest> 是一个 union,所以这个表达式的结果是一个元组的 union,而不再是一个单独的元组,这就节省了一层迭代。

(我后悔了,所有的 元组 概念应该直接写成 tuple 的……)

对于 Record 结构

对于 Record 结构来说,思路是一样的,但这次我们就没这么幸运了。

首先,我们选择 Record 的字段 key 作为 reduce 的维度。通过 keyof T 很容易就能将它们作为一个 union 提取出来。

理想情况下,我们能将这个 keys 的 union 解析为 one 和 rest 两个部分,但现阶段还比较困难。所以我们先降个级,从一个低效但有效的方式开始:

image

如此这般,我们其实对每一个 key 同时发起了算法计算,但实际上我们只需要对其中一个 key 开始就可以了。这一对计算得出的结果是一样的,所以他们的 union 还是同样的结果。最后的结果是正确的,但对于 ts 编译器来说并不高效。但此刻,我们先就这样办,等做完了挑战,再回过头来优化。

继续:

image

如前所述,上图展示的就是两层迭代结构:

  1. 一层是对当前 key 对应的 values 映射过后的值进行迭代
  2. 另一层是对 record 剩余部分进行映射后的结果进行迭代

继续:

image

这个语句将结果 A 和结果 B 合并成最终结果,就像处理 tuple 时的 [Value, ...DistributeTuple<Rest>] 语句一样。

需要注意一下第一个语句:

image

这是一个边界值判断,否则,对空 record 会算出 never 结果。

优化

好了,该回头来拯救我们的 ts 编译器了。按照我们的算法设计,我们需要将 union 解析为 one 和 rest 两部分。下面是我们需要用到的工具:

type UnionToIntersection<T> = (T extends any ? (x: T) => void : never) extends (x: infer U) => void ?  U : never
type PopUnion<T> = UnionToIntersection<T extends any ? (x: T) => void : never> extends (x: infer U) => void ? U : never
type OneAndRest<T> = [PopUnion<T>, Exclude<T, PopUnion<T>>]

我们最终需要使用的是 OneAndRest,它接收 1|2|3 并返回 [3, 1|2]

UnionToIntersection 是我们的老朋友了,参考 Union to Intersection.

这里起到关键作用的新朋友是 PopUnion。它接收一个 union,并返回其中一个单独的类型。

我还不太理解它的工作原理,但根据天才网友们的说法,它依赖了 ts 编译器的内部实现:当 ts 编译器对一个重载函数进行 infer 时,它只会返回这堆函数中的最后一个。

现在我们可以开始优化 DistributeRecord 了。我们把原来的这一行

image

替换为

image

然后对后续代码中的相关引用做一些修改就可以了,并不困难。

优化之后,计算结果和之前是一模一样的,但编译器变得更高效也更高兴了,我们可真是大好人。