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 相關的內容了，整個專案還是比我預想的複雜了一些，但核心邏輯主要還是上文中粗體的部分。

看了半天好像也沒啥特別可以借鑑的地方……