不可变数据流的完美解决方案——Immer 源码解读

1,677 阅读14分钟

更新:现在多了一个状态库适配了 immer:jotai-immer

文章目录

Immer 现在已经成为了很多状态库的标配,包括 Redux, zustand,在去年的库使用频率的调研中,Immer 也是名列前茅。

image.png

为了照顾没有 Immer 使用经验的同学,本文并不直接深入源码,而是包含三个方面:

  1. 基于 React,说明原先的问题,探讨为什么需要 Immer
  2. Immer 的基础用法,为源码解读做铺垫
  3. Immer 源码解读

Immer 官网

顾虑重重的深拷贝

1

在 JS 中,实现深拷贝一般来说有下面几种方式:

  1. 调用库函数,如 lodash 下的 cloneDeep
  2. 使用 JSON.stringifyJSON.parse ,但是这种没办法拷贝不属于 JSON 格式的值,例如 Date 对象、函数、Map、Set 等。
  3. 使用 structuredClone,比第 2 种方法要好,它支持像 Map、Set 这些 JS 内置的对象,但是它不拷贝原型链上的,也不能拷贝函数。

事实上,在实际的场景中我们一般还是使用第 1 种方法。

2

在 React 中,有一个或许可以称之为比较反直觉的现象,可以先看代码,再看下方的 GIF 图:

function App() {
  const [countRef, setCount] = useState({ count: 1 })

  console.log('rerendered')

  function handleClick() {
    countRef.count += 1
    setCount(countRef)
  }

  return (
    <button onClick={handleClick}>
      count is {countRef.count}
    </button>
  )
}

CleanShot 2022-09-25 at 09.40.20.gif

可以看到,当我们点击按钮的时候,值的改变没有表现出来,同时也没有触发组件的重新渲染,也就是说,countRef 的索引一直没变,setCount 触发组件渲染失效了。

最朴素的解决的方法就是每次重新赋值 countRef 之前都深拷贝一份,上面我们列的三个方法可以随便选用一个,在这里我们选择不太常见的第三个:

 function handleClick() {
    const newCountRef = structuredClone(countRef)
    newCountRef.count += 1
    setCount(newCountRef)
  }

此时再看:

CleanShot 2022-09-25 at 09.48.09.gif

这种解决方案可以说是简单粗暴,我们在某些时候也是懒得想,直接使用。但它也有其局限性:它在数据量大的时候就会出现性能瓶颈。

如果我们一个列表有非常多项,如果我们只更改了列表中某一项的一个值,接着对一整个列表做一遍深拷贝的话,可想而知,频率很高的深拷贝也是一个比较耗费性能的操作(具体情况需要分析得知)。

临时救火的浅拷贝

深拷贝既然不行,我们可能得采用浅拷贝了。

于是,最简单的方法可能形如下面这样:

list[1].title = 'new title'
setList([...list])

我们利用浅拷贝把 list 的地址更改一下,这样也能触发更新。在 React 中,如果我们的子组件没有使用 React.memo 包裹,父组件有了更新会触发下面每一个子组件重新渲染,除非子项是一个 Text Node 节点,就算是子组件没有接受任何 props,新旧 props 对象都是空对象 {},进行比较也不等,如下图所示:

CleanShot 2022-09-25 at 14.43.30@2x.png

好处是利用 React 的这个特性,不深拷贝对象,只改变列表对象的索引,就能触发列表的更新。坏处便是为了造成了很多子项的 Fiber Diff,为了解决这个问题,我们很多时候都得包裹 React.memo。而正是由于使用 React.memo 包裹,也引发了后面的问题。

我们可以设计下面这个例子,你可以在线预览调试一下,推荐点击右上角【查看详情】去查看,当点击 【Change Title】按钮,没有任何反应,便是我们上面说的问题了:

如果你看完了上面的例子,应该就能体会出什么问题了,如何去解决这个问题呢?最朴素的方法则是在自定义 React.memocompare 函数:

const Item = React.memo(({ subject }: { subject: any }) => {
  return <div>{subject.title}</div>
}, (prev, cur) => {
  if (prev.subject !== cur.subject) {
    return false
  }
  return true
})

但是太遗憾了,还是不行,因为我们更改的 list 和新的是同一个索引地址,也就是说 prevcur 都是同一个。你可以在上面的在线代码片段里修改试一下。

难道这个问题就无解了吗?当然不是的,那就是不传递对象下来,而是把值解构出来传下来,这个方法是用到了 memo 函数默认的比较函数是比较 props 每一项的特点:

const Item = React.memo(({ title }: { title: string }) => {
  return <div>{title}</div>
})

又或者改值的时候:

  function handleChange() {
    setList([
      { title: 'Javascript changed', name: 'foo' },
      list[1],
      list[2]  
    ])
  }

相信,使用 React 的很多同学都会采用过上面两种写法,甚至还有更麻烦的:

setObject({
    ...level1,
    appendLevel1: {
        ...appendLevel1,
        appendLevel2: {
            ...
        }
    }
})

上述代码在以前的 Redux 中可谓遍地都是,笔者之前的团队曾在小程序中引入了 Redux,但是没有使用 Immer。写起来麻烦是一个方面,性能也不见得好到哪里去,由于小程序性能有限,在列表达到 200 项的时候,一些低端机器就开始闪退了。

Immer 如何解决

Immer 的出现解决了上面的问题,并且完美的兼容 React.memo,上面的例子只需改造成:

import produce from "immer"

// ...

function handleChange() {
    const nextList = produce(list, (draft) => {
      draft[0].title = 'Javascript Changed'
    })

    setList(nextList)
}

在 Hook 中的使用也非常方便:

import React, { useCallback } from "react";
import { useImmer } from "use-immer";

function App() {
    const [todos, setTodos] = useImmer([
        {
            id: "React",
            title: "Learn React",
            done: true
        },
        {
            id: "Immer",
            title: "Try Immer",
            done: false
        }
    ]);

    const handleToggle = useCallback((id) => {
        setTodos((draft) => {
            const todo = draft.find((todo) => todo.id === id);
            todo.done = !todo.done;
        });
    }, []);

    // 其余部分省略...
}

如果对完整示例感兴趣,请参阅 官网

Redux Toolkit 也是内置继承了 Immer,主要的改动在 Reducer ,改为了下面这样的格式书写:

export const counterSlice = createSlice({
  name: 'counter',
  initialState,
  reducers: {
    increment: (state) => {
      state.value += 1
    },
    decrement: (state) => {
      state.value -= 1
    },
  },
})

说到这里,不敢相信谁能抵挡得住,想去了解这个库源码的诱惑。话不多说,现在就为大家带来它的源码解读。

本文只解析 produce 函数的逻辑。

Immer 源码解读

前置知识

Proxy

  1. Proxy.revocable 阅读资料:我们可以调用 new Proxy() 去创建一个代理,也可以使用 Proxy.revocable,但是后者可以销毁掉 proxy 对象。
var revocable = Proxy.revocable({}, {
  get(target, name) {
    return "[[" + name + "]]";
  }
});
var proxy = revocable.proxy;
proxy.foo;              // "[[foo]]"

revocable.revoke(); // proxy 不起作用了,再取值会报错

树的深度优先遍历

对象 的整个代理过程就是一个深度优先遍历,伪代码如下:

function traverse (obj) {
    for (let key of Object.keys(obj)) {
        doSomethingIn(obj[key])
        traverse(obj[key])
    }
}

接下来,我们把流程拆解为 7 步。

Step 1

预备工作,先克隆下代码,安装依赖:

git clone https://github.com/immerjs/immer.git
pnpm i

这个 是入口函数的地址,如果你想和我一块,可以找到这个位置。

如果你想调试代码,可以直接起一个安装了 Immer 的项目进行 Debugger,但是这样某些变量名会被转译,有点妨碍阅读,下面是一个很笨但是管用的方法:

  1. Fork 项目
  2. 在根目录新建一个 Vite 环境
  3. 引入入口文件导出的 produce 函数
  4. 有报错,解决它

Step 2

我们先给一个最小的实现,让大家快速抓住脉络。我们先来说一下它的实现思路。

简单的说,Immer 把原始对象的每一个对象里的值扩充为:

var obj = {
    a: a的值
}

// 转化为
var proxyTarget = {
    a: {
        base_: a的值
        copy_: a的深拷贝
    }
}

也就是说,本来的值存为 base_, 保留一份对原始对象的深拷贝(并不准确,只是为了大家理解),记作 copy_。同时会维护一个 proxy 对象,这个 proxy 对象就是 recipe 函数的 draft 对象。

当我们对对象进行修改,Immer 会把这个修改代理到 copy_ 上面。最后再把 copy_ 返回(并不准确,只是为了大家理解)。

有一点值得注意,在返回最终结果的时候并不是简单的返回 copy_ 对象,如果当前值没有被修改,那就返回 base_ 对应的值。

我们可以把对象想象成一棵树,当某一个孩子节点被修改了,它及其它的 parent 节点都会被标记为被修改,此时返回 copy_ 对象上的值;而除去这个链路上的节点(它的 sibling 节点,它 parent 的 sibling 节点)都直接返回 base_ 对象上的值。

所以整体流程可以是:

  1. 收集更改
  2. 合并更改

为了简单起见,本文只考虑对象的场景,数组的情况也不进行考虑,下面不再做特别说明。

Step 3

首先,我们给出 produce 函数的主体逻辑和注释,

produce = (base: any, recipe: any) => {
    let result;

    // 判断传入的值是不是合法的对象
    if (isDraftable(base)) {
        // 创建一个上下文,后面根据上下文可以还原出值
        const scope = enterScope(this)
        
        // 对 base 对象创建代理,以便劫持所有的修改
        const proxy = createProxy(this, base, undefined)

        // 调用修改函数,我们 recipe 可能只操作 proxy,不返回值
        result = recipe(proxy)

        // 合并出最终的结果,返回
        return processResult(result, scope)
    }
}

接下来我们便逐个方法进行解析。

Step 4

if (isDraftable(base)) {
    ...
}

如果是对象,我们才可以继续,代码可以简化为:

function isDraftable(value: any) {
    if (!value || typeof value !== "object") return false
    
    return true
}

Step 5

const scope = enterScope(this)

目前我们还没有用到它,大家看一下实现即可,可以略过,等用到再回来看,实现如下:

let currentScope: ImmerScope | undefined

export function enterScope(immer: Immer) {
    return (currentScope = createScope(currentScope, immer))
}

function createScope(parent_, immer_): ImmerScope {
    return {
        drafts_: [],
        parent_,
        immer_,
    }
}

Step 6

const proxy = createProxy(this, base, undefined)

代码如下:

export function createProxy<T extends Objectish>(
    immer: Immer,
    value: T,
    parent?: ImmerState
): Drafted<T, ImmerState> {
    const draft: Drafted = createProxyProxy(value, parent)

    const scope = parent ? parent.scope_ : getCurrentScope()
    scope.drafts_.push(draft)
    return draft
}

逻辑大致如下:

  1. 调用 createProxyProxy 绑定好代理关系,返回的是一个 Proxy 对象
  2. scope 入当前 draft 对象,用途后面解析
  3. 返回 Proxy 对象

核心是 createProxyProxy,它是收集更改过程中灵魂的函数。

export function createProxyProxy<T extends Objectish>(
	base: T,
	parent?: ImmerState
): Drafted<T, ProxyState> {
    const state: ProxyState = {
        scope_: parent ? parent.scope_ : getCurrentScope()!,
        modified_: false,
        finalized_: false,
        assigned_: {},
        parent_: parent,
        base_: base,
        draft_: null as any, // set below
        copy_: null,
        revoke_: null as any,
        isManual_: false
    }

    let target: T = state as any
    let traps: ProxyHandler<object | Array<any>> = objectTraps
    const {revoke, proxy} = Proxy.revocable(target, traps)
    state.draft_ = proxy as any
    state.revoke_ = revoke
    return proxy as any
}

流程是:

  1. 定义好 Proxy 的 target 对象,即为 state,所有的重要信息都存在这个对象里
  2. ojectTraps 去拦截对此对象的修改
  3. 返回此 proxy

state 的每一个属性穿插后续流程,我们给出值的具体含义:

变量名含义
modified标记是否被修改,如果修改,返回数据的时候从 copy_ 字段取
finalized标记是否已经完成修改,已完成修改对 proxy 的操作将不被接受
parent它的上层 proxy 对象
base_原始数据
draft_原始数据对应的 proxy 对象
copy_原始数据的一份拷贝,最初是浅拷贝
revoke_销毁 proxy 对象用

接下来是 objectTraps 的实现,我们主要看它的 get 方法和 set 方法,关键点是 按需绑定 proxy,只有在读到那个值的时候才绑定。

几个重要的工具函数:

// 总是取最新的值,有 copy_ 就不取 base_
// 这样能取到最新修改的值
export function latest(state: ImmerState): any {
    return state.copy_ || state.base_
}

// DRAFT_STATE 是一个内置的 Symbol 对象
// 如果 draft 是原始对象,肯定没有 DRAFT_STATE,返回本身的值
// 如果 state 有值,说明 draft 是 proxy,那调用 latest(state)[prop] 取值
// DRAFT_STATE 其实就是在 proxy 的 get 里代理了的,返回它本身,详情见 get 函数第一行
function peek(draft: Drafted, prop: PropertyKey) {
    const state = draft[DRAFT_STATE]
    const source = state ? latest(state) : draft
    return source[prop]
}

// 把 `base_` 上的每一个自有属性移动到 `copy_` 上
// 做的浅拷贝,也就是说,在此刻:
// state.copy_ !== state.base_ 
// 但是 state.copy_[someKey] ===  state.base_[someKey]
export function prepareCopy(state: {base_: any; copy_: any}) {
    if (!state.copy_) {
            state.copy_ = shallowCopy(state.base_)
    }
}

一定要注意,上面三个函数绝对要理解用途,它是理解后面 set、get 的核心。不懂的话可以多读几遍。

在上面的基础上,我们再来看 get 函数:

// https://github.com/mysteryven/immer/blob/d91a6597e92570086b329ba5b197c18d211077db/src/core/proxy.ts#L101
get(state, prop) {
    if (prop === DRAFT_STATE) return state

    // 取最新的值,此时 value 可能是 proxy 对象,也可能是原始对象
    const source = latest(state)
    const value = source[prop]

    // 修改结束或者不是对象结构,就直接返回
    // state.finalized_ 在后面合并依赖时被设置为 false
    // 也是调用完 recipe 函数之后
    if (state.finalized_ || !isDraftable(value)) {
        return value
    }

    // 看看二者是否相等,右边值一定是原始值,相等说明还没有被代理
    if (value === peek(state.base_, prop)) {
        // 浅拷贝 `base_` 的属性 到 `copy_`
        prepareCopy(state)
        // `copy_` 的 prop 这一项 重新赋值为 proxy 对象
        return (state.copy_![prop as any] = createProxy(
                state.scope_.immer_,
                value,
                state
        ))
    }
    return value
},

可能有同学看到了这里会疑问,为什么在 prepareCopy 浅拷贝到了 copy_, 又接着把 copy_ 设置为 Proxy 对象,这不是多此一举吗?

原因有两个:

  1. 我们 base_ 的原始数据可能不是可写的(writeable),我们得先把它拷贝到 copy_ 上,同时把它变成可写的。
  2. 我们的值最终要写入一个对象中,它的值也承接了一个桥梁作用,而且我们对所有属性都做了浅拷贝,但是只 prop 这一项是做了 Proxy 代理

接着看 set 函数,理解此函数最关键的一步就是要知道 copy_ 的数据格式,它是拥有和 base_ 一样 key 值,但是 value 可能是原始值,也可能是 proxy 对象。

我们来看一下它的函数主体:

set(
    state: ProxyObjectState,
    prop: string
    value
) {
    // 没有被标记为修改,说明第一次修改
    // 可能需要进行拷贝、标记修改的操作
    if (!state.modified_) {
        // 拿到最新的值
        const current = peek(latest(state), prop)
        //如果是 proxy,拿它本身的 state 对象
        const currentState: ProxyObjectState = current?.[DRAFT_STATE];
        
        // 如果 value 和 base_ 相等,那没必要赋值
        if (currentState && currentState.base_ === value) {
                state.copy_![prop] = value
                state.assigned_[prop] = false
                return true
        }
        
        // 浅拷贝 state 到 `copy_` 中
        // 如 a.b.c.d.e = {x: 'hello world'}
        // d.e 算是赋值了
        // get 那里赋值 `copy_` 只到 d 这一层
        // e 这一层还是没有被拷贝。
        prepareCopy(state)
        
        // 标记 state 和 state 的 parent 为 modifed_
        markChanged(state)
    }

    // 直接赋值 copy_ 的值为 value
    // 此时不再是 proxy 对象了
    state.copy_![prop] = value
    state.assigned_[prop] = true
    
    return true
},


export function markChanged(state: ImmerState) {
    if (!state.modified_) {
        state.modified_ = true
        if (state.parent_) {
                markChanged(state.parent_)
        }
    }
}

Step 7

经过了 Step 6, 也就是调用完了 producerecipe 函数,调用了 get 或者 set,最终会进入合并阶段,即 Step 3 的 processResult 函数。

这里主要是根据 recipe 函数返回值做区分:

export function processResult(result: any, scope: ImmerScope) {
    const baseDraft = scope.drafts_![0]
    const isReplaced = result !== undefined && result !== baseDraft

    // 有返回值做替换
    if (isReplaced) {
        if (baseDraft[DRAFT_STATE].modified_) {
            revokeScope(scope)
        }
        if (isDraftable(result)) {
            result = finalize(scope, result)
            if (!scope.parent_) maybeFreeze(scope, result)
        }
    } else {
        result = finalize(scope, baseDraft, [])
    }

    revokeScope(scope)

    return result
}

特别说明,没有返回值的情况,也就是 recipe 函数这样调用:

state = produce(state, draft => {
    draft.name = "Michel"
})

无论如何,关键的就两个函数:finalizerevokeScope

  1. finalize 将根据 base_copy_ 合并出最终的结果
  2. revokeScope 注销掉一个个 proxy 对象

理解 finalize 函数的关键还是理解 copy_base_ 值关系:

function finalize(rootScope: ImmerScope, value: any) {
    const state: ImmerState = value[DRAFT_STATE]

    // 说明是原始对象,直接遍历它的每一个 child
    if (!state) {
        each(
            value,
            (key, childValue) =>
                // 作用是为每一个属性也调用 finalize
                finalizeProperty(rootScope, state, value, key, childValue),
        )
        return value
    }

    // 没有修改过,直接返回
    if (!state.modified_) {
        return state.base_
    }

    // 是 state 对象格式,并且被修改过了
    if (!state.finalized_) {
        // 先标记为修改完成,这时候再进行 set 就不生效了
        state.finalized_ = true
        
        // copy_ 的值可能是 proxy 对象,可能是原始对象,可能是原生值
        // 最终要把 copy_ 的值为 proxy 对象的变为原始对象,其他的保持不变。
        const result = state.copy_
        each(
            result,
            (key, childValue) =>
                finalizeProperty(rootScope, state, result, key, childValue, path)
        )
    }
    return state.copy_
}

export function isDraft(value: any): boolean {
    return !!value && !!value[DRAFT_STATE]
}

function finalizeProperty(
    rootScope: ImmerScope,
    parentState: undefined | ImmerState,
    targetObject: any,
    prop: string | number,
    childValue: any,
) {
    // 是不是 Proxy 对象
    if (isDraft(childValue)) {
        // Drafts owned by `scope` are finalized here.
        // 递归的判断内部的逻辑,最终会返回原始对象
        const res = finalize(rootScope, childValue)
        // proxy 对象重写成原始对象
        targetObject[prop] = res
    } else return
}

理解的关键是如果把根级别的 copy_ 看做树,它的属性的值可能是 proxy 对象,可能是原始值(可能是新的或者和 base_ 一样),但是树的叶子节点肯定是原始值。甚至,叶子节点往上也可能是原始值而不是 proxy 对象,这就保证了递归过程有终止。而遍历方式是先遍历孩子(const res = finalize(rootScope, childValue))再赋值本身( targetObject[prop] = res),就保证了赋值的时候,值一定是原始值了,最终拿到了结果 copy_ 便是正确的。

revokeScope 和主流程关系不大,但是为了保证完整性,还是给出所有的函数:

export function revokeScope(scope: ImmerScope) {
    leaveScope(scope)
    // 销毁每一个 drafts 里的 proxy 对象
    scope.drafts_.forEach(revokeDraft)
    scope.drafts_ = null
}

export function leaveScope(scope: ImmerScope) {
    if (scope === currentScope) {
            currentScope = scope.parent_
    }
}

function revokeDraft(draft: Drafted) {
    const state: ImmerState = draft[DRAFT_STATE]
    if (
        state.type_ === ProxyType.ProxyObject ||
        state.type_ === ProxyType.ProxyArray
    )
        state.revoke_()
    else state.revoked_ = true
}

到这里,我们的 Immer 源码解读就完成了。

总结

Immer 借助 Proxy 的能力,对我们读值、取值进行代理,并且巧妙的在 base_copy_ 进行中转,整个过程都是按需修改,如果我们有了更改,才会按需的增加 Proxy 去监听,最后再把二者值合并出最终的对象。