Immer 全解析

Example

import produce from "immer"

const baseState = [
    {
        todo: "Learn typescript",
        done: true
    },
    {
        todo: "Try immer",
        done: false
    }
]

// baseState 不变，nextState 是变更后的新对象
const nextState = produce(baseState, draftState => {
    draftState.push({todo: "Tweet about it"})
    draftState[1].done = true
})
复制代码

初识 Immer

第一次听说 Immer 差不多在几个月前吧，那会儿写了个状态管理库想在公司推广，组内同学发了 Immer 的 GitHub 地址给我，说是有个基于 Proxy 的状态管理库，自称性能很好，我们能用上么？我就去瞄了几眼，回复说这称不上状态管理库吧，概念上更贴近 Immutable.js，就是用来方便操作 immutable 数据的。它提供给用户一个 draftState，用户可以随意对它进行修改，最后会返回一个新数据，原数据不变。当时也稍微看了下它的核心原理，draftState 是个 Proxy，对它的读写操作会走到内部定义好的 getter/setter 里，简单来说就是当你获取 draftState 内部的对象时，它都会返回一个 Proxy，而当你进行赋值时，它都会对原对象的 copy 对象进行赋值。最后返回 copy 对象。

我的项目里其实也用到了 Proxy，是用来简化消息发送相关操作的。不过我项目里的状态是 mutable 的（一开始其实是 immutable 的，后来为了实现某些功能，事情就发生了变化……），所以 Immer 与我来说有些鸡肋，就没怎么放在心上。

源码解析

然而换了公司之后，又频频从新同事口中听到 Immer，想用到我们的项目里。虽然我觉得它还是不符合我们的场景，并不打算用，但听得多了就觉得还是完整地看一下源码吧，或许能借鉴点什么边边角角的东西呢……

produce

produce 是直接暴露给用户使用的函数，它是 Immer 类的一个实例方法（可以先不看代码直接看我下面的解释）：

export class Immer {
    constructor(config) {
        assign(this, configDefaults, config)
        this.setUseProxies(this.useProxies)
        this.produce = this.produce.bind(this)
    }
    produce(base, recipe, patchListener) {
        // curried invocation
        if (typeof base === "function" && typeof recipe !== "function") {
            const defaultBase = recipe
            recipe = base

            // prettier-ignore
            return (base = defaultBase, ...args) =>
                this.produce(base, draft => recipe.call(draft, draft, ...args))
        }

        // prettier-ignore
        {
            if (typeof recipe !== "function") throw new Error("if first argument is not a function, the second argument to produce should be a function")
            if (patchListener !== undefined && typeof patchListener !== "function") throw new Error("the third argument of a producer should not be set or a function")
        }

        let result

        // Only plain objects, arrays, and "immerable classes" are drafted.
        if (isDraftable(base)) {
            const scope = ImmerScope.enter()
            const proxy = this.createProxy(base)
            let hasError = true
            try {
                result = recipe.call(proxy, proxy)
                hasError = false
            } finally {
                // finally instead of catch + rethrow better preserves original stack
                if (hasError) scope.revoke()
                else scope.leave()
            }
            if (result instanceof Promise) {
                return result.then(
                    result => {
                        scope.usePatches(patchListener)
                        return this.processResult(result, scope)
                    },
                    error => {
                        scope.revoke()
                        throw error
                    }
                )
            }
            scope.usePatches(patchListener)
            return this.processResult(result, scope)
        } else {
            result = recipe(base)
            if (result === undefined) return base
            return result !== NOTHING ? result : undefined
        }
    }
复制代码

produce 接收三个参数，正常来说 base 是原数据，recipe 是用户执行修改逻辑的地方，patchListener 是用户接收 patch 数据然后做一些自定义操作的地方。

produce 一开始的逻辑看注释是为了柯里化（其实并不是严格的柯里化，不过和本文内容无关，略过不谈），它判断了下 base 是不是函数，如果是的话把 base 赋值给 recipe，然后再返回一个接收 base 的函数，什么意思呢？就是一般情况你是像produce(base, (draft) => { ... })这样调用 produce，但如果某些情况下你要先接收 recipe 函数再接收 base，那你可以像produce((draft) => { ... })(base)这样调用，最常见的场景是配合 React 的 setState：

// state = { user: { age: 18 } }
this.setState(
    produce(draft => {
        draft.user.age += 1
    })
)
复制代码

当然你也可以传入默认 base，const changeFn = produce(recipe, base)，可以直接changeFn()也可以changeFn(newBase)，newBase 会覆盖之前的 base。

接下来是主流程：

• 如果 base 是对象（包括数组），能生成 draft，则：
  • 执行const scope = ImmerScope.enter()，生成一个 ImmerScope 的实例 scope，scope 和当前的 produce 调用绑定
  • 执行this.createProxy(base)创建 proxy（draft），并执行scope.drafts.push(proxy)将 proxy 保存到 scope 里
  • 以 proxy 为参数调用用户传入的 recipe 函数，并把返回值保存为 result
  • 如果执行 recipe 期间没有出错则调用scope.leave，把 ImmerScope.current 重置为初始状态（这里是 null），如果出错了则执行scope.revoke()，重置所有状态。
  • 判断 result 是否为 promise，是则返回result.then(result => this.processResult(result, scope))，否则直接返回this.processResult(result, scope)（返回前其实还要执行scope.usePatches(patchListener)，patch 相关的东西不算主流程，先不管）
• 如果 base 不能生成 draft，则：
  • 执行result = recipe(base)
  • result 为 undefined 直接返回 base；否则判断 result 是否为NOTHING（一个内部标记），是则返回 undefined，否则返回 result

整个 produce 主要就做了三个事情：

• 调用createProxy生成 draft 供用户使用
• 执行用户传入的 recipe，拦截读写操作，走到 proxy 内部的 getter/setter
• 调用processResult解析组装最后的结果返回给用户

接下来我们一步步探究涉及到的部分。

创建 draft

你会发现 Immer 的 class 声明里并没有 createProxy 这个实例方法，但却能在 produce 内执行this.createProxy(base)。Is it magic？ 实际上 createProxy 是存在于 proxy.js 和 es5.js 文件内的，es5.js 里的内容是个兼容方案，用于不支持 Proxy 的环境，immer.js 的开头会 import 两个文件的内容：

import * as legacyProxy from "./es5"
import * as modernProxy from "./proxy"
复制代码

在 Immer 的 constructor 里会执行this.setUseProxies(this.useProxies)，useProxies 用来表示当前环境是否支持 Proxy，setUseProxies 里会判断 useProxies：

• is true：assign(this, modernProxy)
• is false: assign(this, legacyProxy)

这样createProxy函数就被挂载到this上了，这里我们详细看看 proxy.js 里的createProxy：

export function createProxy(base, parent) {
    const scope = parent ? parent.scope : ImmerScope.current
    const state = {
        // Track which produce call this is associated with.
        scope,
        // True for both shallow and deep changes.
        modified: false,
        // Used during finalization.
        finalized: false,
        // Track which properties have been assigned (true) or deleted (false).
        assigned: {},
        // The parent draft state.
        parent,
        // The base state.
        base,
        // The base proxy.
        draft: null,
        // Any property proxies.
        drafts: {},
        // The base copy with any updated values.
        copy: null,
        // Called by the `produce` function.
        revoke: null
    }

    const {revoke, proxy} = Array.isArray(base)
        ? // [state] is used for arrays, to make sure the proxy is array-ish and not violate invariants,
          // although state itself is an object
          Proxy.revocable([state], arrayTraps)
        : Proxy.revocable(state, objectTraps)

    state.draft = proxy
    state.revoke = revoke

    scope.drafts.push(proxy)
    return proxy
}
复制代码

• 根据 base 构建一个 state 对象，里面的属性我们等用到的时候再细说
• 判断 base 是否为数组，是则基于 arrayTraps 创建[state]的 Proxy，否则基于 objectTraps 创建state的 Proxy

arrayTraps 基本就是转发参数到 objectTraps，而 objectTraps 里比较关键的是 get 和 set，对 proxy 的取值和赋值操作都会被这两个函数拦截。

拦截取值操作

function get(state, prop) {
    if (prop === DRAFT_STATE) return state
    let {drafts} = state

    // Check for existing draft in unmodified state.
    if (!state.modified && has(drafts, prop)) {
        return drafts[prop]
    }

    const value = source(state)[prop]
    if (state.finalized || !isDraftable(value)) return value

    // Check for existing draft in modified state.
    if (state.modified) {
        // Assigned values are never drafted. This catches any drafts we created, too.
        if (value !== state.base[prop]) return value
        // Store drafts on the copy (when one exists).
        drafts = state.copy
    }

    return (drafts[prop] = createProxy(value, state))
}
复制代码

get 接收两个参数，第一个为 state，即创建 Proxy 时传入的第一个参数（目标对象），第二个参数为 prop，即想要获取的属性名，具体逻辑如下：

• 若 prop 为DRAFT_STATE则直接返回 state 对象（会在最后处理结果时用到）
• 取 state 的 drafts 属性。drafts 中保存了state.base子对象的 proxy，譬如base = { key1: obj1, key2: obj2 }，则drafts = { key1: proxyOfObj1, key2: proxyOfObj2 }
• 若 state 尚未被修改并且 drafts 中存在 prop 对应的 proxy，则返回该 proxy
• 若state.copy存在，则取state.copy[prop]，否则取state.base[prop]，存于 value
• 若 state 已经结束计算了或者 value 不能用来生成 proxy，则直接返回 value
• 若 state 已被标记修改
  • 若value !== state.base[prop]则直接返回 value
  • 否则把state.copy赋值给 drafts（copy 里也包含了子对象的 proxy，具体会在 set 部分细说）
• 若未提前返回则执行createProxy(value, state)生成以 value 为 base、state 为 parent 的子 state 的 proxy，存到 drafts 里并返回

讲完了 get，我们发现它就是用来生成子对象的 proxy，缓存 proxy，然后返回 proxy，如果不能生成 proxy 则直接返回一个值。

拦截赋值操作

function set(state, prop, value) {
    if (!state.modified) {
        // Optimize based on value's truthiness. Truthy values are guaranteed to
        // never be undefined, so we can avoid the `in` operator. Lastly, truthy
        // values may be drafts, but falsy values are never drafts.
        const isUnchanged = value
            ? is(state.base[prop], value) || value === state.drafts[prop]
            : is(state.base[prop], value) && prop in state.base
        if (isUnchanged) return true
        markChanged(state)
    }
    state.assigned[prop] = true
    state.copy[prop] = value
    return true
}
复制代码

set 接受三个参数，前两个和 get 的一样，第三个 value 是将要赋予的新值，具体逻辑如下：

• 先判断 state 是否被标记更改，若没有，则：

  • 判断新值和旧值是否相等，若相等则直接返回，啥都不做
  • 否则执行markChanged(state)（后面细讲）

• 将state.assigned[prop]置为 true，标记该属性被赋值

• 将 value 赋值给state.copy[prop]

整个 set 的核心其实是标记修改并把新值赋给 copy 对象的对应属性，现在我们看下 margeChanged:

function markChanged(state) {
    if (!state.modified) {
        state.modified = true
        state.copy = assign(shallowCopy(state.base), state.drafts)
        state.drafts = null
        if (state.parent) markChanged(state.parent)
    }
}
复制代码

一个 state 只需被标记一次，具体如下：

• 把state.modified置为 true
• 浅拷贝state.base，并把state.draftsassign 到拷贝对象，赋值给state.copy。也就是说state.copy中含有子对象的 proxy，会在 get 中用到，之前我们已经说过了
• 把state.drafts置为 null
• 如果 state 有 parent，递归执行markChanged(state.parent)。这很好理解，譬如draft.person.name = 'Sheepy'这个操作，我们不止要把 person 标记修改，也要把 draft 标记修改

解析结果返回

processResult(result, scope) {
  const baseDraft = scope.drafts[0]
  const isReplaced = result !== undefined && result !== baseDraft
  this.willFinalize(scope, result, isReplaced)
  if (isReplaced) {
    if (baseDraft[DRAFT_STATE].modified) {
      scope.revoke()
      throw new Error("An immer producer returned a new value *and* modified its draft. Either return a new value *or* modify the draft.") // prettier-ignore
    }
    if (isDraftable(result)) {
      // Finalize the result in case it contains (or is) a subset of the draft.
      result = this.finalize(result, null, scope)
    }
    if (scope.patches) {
      scope.patches.push({
        op: "replace",
        path: [],
        value: result
      })
      scope.inversePatches.push({
        op: "replace",
        path: [],
        value: baseDraft[DRAFT_STATE].base
      })
    }
  } else {
    // Finalize the base draft.
    result = this.finalize(baseDraft, [], scope)
  }
  scope.revoke()
  if (scope.patches) {
    scope.patchListener(scope.patches, scope.inversePatches)
  }
  return result !== NOTHING ? result : undefined
}

复制代码

虽然 Immer 的 Example 里都是建议用户在 recipe 里直接修改 draft，但用户也可以选择在 recipe 最后返回一个 result，不过得注意“修改 draft”和“返回新值”这个两个操作只能任选其一，同时做了的话processResult函数就会抛出错误。我们重点关注直接操作 draft 的情况，核心逻辑是执行result = this.finalize(baseDraft, [], scope)，返回 result 的情况也是相似的，都要调用finalize，我们看一下这个函数：

/**
 * @internal
 * Finalize a draft, returning either the unmodified base state or a modified
 * copy of the base state.
 */
finalize(draft, path, scope) {
  const state = draft[DRAFT_STATE]
  if (!state) {
    if (Object.isFrozen(draft)) return draft
    return this.finalizeTree(draft, null, scope)
  }
  // Never finalize drafts owned by another scope.
  if (state.scope !== scope) {
    return draft
  }
  if (!state.modified) {
    return state.base
  }
  if (!state.finalized) {
    state.finalized = true
    this.finalizeTree(state.draft, path, scope)

    if (this.onDelete) {
      // The `assigned` object is unreliable with ES5 drafts.
      if (this.useProxies) {
        const {assigned} = state
        for (const prop in assigned) {
          if (!assigned[prop]) this.onDelete(state, prop)
        }
      } else {
        const {base, copy} = state
        each(base, prop => {
          if (!has(copy, prop)) this.onDelete(state, prop)
        })
      }
    }
    if (this.onCopy) {
      this.onCopy(state)
    }

    // At this point, all descendants of `state.copy` have been finalized,
    // so we can be sure that `scope.canAutoFreeze` is accurate.
    if (this.autoFreeze && scope.canAutoFreeze) {
      Object.freeze(state.copy)
    }

    if (path && scope.patches) {
      generatePatches(
        state,
        path,
        scope.patches,
        scope.inversePatches
      )
    }
  }
  return state.copy
}
复制代码

我们略过类似钩子函数的onDelete和onCopy，只看主流程：

• 通过 draft 拿到 state（在 createProxy 里生成的 state 对象，包含 base、copy、drafts 等属性）
• 若 state 未被标记修改，直接返回state.base
• 若 state 未被标记结束，执行this.finalizeTree(state.draft, path, scope，最后返回state.copy

我们看下finalizeTree：

finalizeTree(root, rootPath, scope) {
  const state = root[DRAFT_STATE]
  if (state) {
    if (!this.useProxies) {
      state.finalizing = true
      state.copy = shallowCopy(state.draft, true)
      state.finalizing = false
    }
    root = state.copy
  }

  const needPatches = !!rootPath && !!scope.patches
  const finalizeProperty = (prop, value, parent) => {
    if (value === parent) {
      throw Error("Immer forbids circular references")
    }

    // In the `finalizeTree` method, only the `root` object may be a draft.
    const isDraftProp = !!state && parent === root

    if (isDraft(value)) {
      const path =
        isDraftProp && needPatches && !state.assigned[prop]
        ? rootPath.concat(prop)
        : null

      // Drafts owned by `scope` are finalized here.
      value = this.finalize(value, path, scope)

      // Drafts from another scope must prevent auto-freezing.
      if (isDraft(value)) {
        scope.canAutoFreeze = false
      }

      // Preserve non-enumerable properties.
      if (Array.isArray(parent) || isEnumerable(parent, prop)) {
        parent[prop] = value
      } else {
        Object.defineProperty(parent, prop, {value})
      }

      // Unchanged drafts are never passed to the `onAssign` hook.
      if (isDraftProp && value === state.base[prop]) return
    }
    // Unchanged draft properties are ignored.
    else if (isDraftProp && is(value, state.base[prop])) {
      return
    }
    // Search new objects for unfinalized drafts. Frozen objects should never contain drafts.
    else if (isDraftable(value) && !Object.isFrozen(value)) {
      each(value, finalizeProperty)
    }

    if (isDraftProp && this.onAssign) {
      this.onAssign(state, prop, value)
    }
  }

  each(root, finalizeProperty)
  return root
}
复制代码

函数一开始把state.copy赋值给root，最后执行each(root, finalizeProperty)，即以 root 的属性名（prop）和属性值（value）为参数循环调用finalizeProperty，finalizeProperty虽然看着代码很多，实际上就是把 copy 中的 draft(proxy) 属性值替换成draft[DRAFT_STATE].copy（这些 proxy 是在 markChanged 时 assign 上去的，前面我们说过），这样我们就得到了一个真正的 copy，最后可以返回给用户。

总结

由于篇幅问题，就不细讲 patches 相关的内容了，整个项目还是比我预想的复杂了一些，但核心逻辑主要还是上文中粗体的部分。

看了半天好像也没啥特别可以借鉴的地方……