Immer 全解析

7,554 阅读10分钟

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
}

我们略过类似钩子函数的onDeleteonCopy,只看主流程:

  • 通过 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)为参数循环调用finalizePropertyfinalizeProperty虽然看着代码很多,实际上就是把 copy 中的 draft(proxy) 属性值替换成draft[DRAFT_STATE].copy(这些 proxy 是在 markChanged 时 assign 上去的,前面我们说过),这样我们就得到了一个真正的 copy,最后可以返回给用户。

总结

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

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