zustand 源码

zustand 是一个非常时髦的状态管理库，也是 2021 年 Star 增长最快的 React 状态管理库。它的理念非常函数式，API 设计的很优雅，值得学习。

概述

首先介绍 zustand 的使用方法。

创建 store

通过 create 函数创建 store，回调可拿到 get set 就类似 Redux 的 getState 与 setState，可以获取 store 瞬时值与修改 store。返回一个 hook 可以在 React 组件中访问 store。

import create from 'zustand'

const useStore = create((set, get) => ({
  bears: 0,
  increasePopulation: () => set(state => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 })
}))

上面例子是全局唯一的 store，也可以通过 createContext 方式创建多实例 store，结合 Provider 使用：

import create from 'zustand'
import createContext from 'zustand/context'

const { Provider, useStore } = createContext()

const createStore = () => create(...)

const App = () => (
  <Provider createStore={createStore}>
    ...
  </Provider>
)

访问 store

通过 useStore 在组件中访问 store。与 redux 不同的是，无论普通数据还是函数都可以存在 store 里，且函数也通过 selector 语法获取。因为函数引用不可变，所以实际上下面第二个例子不会引发重渲染：

function BearCounter() {
  const bears = useStore(state => state.bears)
  return <h1>{bears} around here ...</h1>
}

function Controls() {
  const increasePopulation = useStore(state => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

如果嫌访问变量需要调用多次 useStore 麻烦，可以自定义 compare 函数返回一个对象：

const { nuts, honey } = useStore(state => ({ nuts: state.nuts, honey: state.honey }), shallow)

细粒度 memo

利用 useCallback 甚至可以跳过普通 compare，而仅关心外部 id 值的变化，如：

const fruit = useStore(useCallback(state => state.fruits[id], [id]))

原理是 id 变化时，useCallback 返回值才会变化，而 useCallback 返回值如果不变，useStore 的 compare 函数引用对比就会为 true，非常巧妙。

set 合并与覆盖

set 函数第二个参数默认为 false，即合并值而非覆盖整个 store，所以可以利用这个特性清空 store：

const useStore = create(set => ({
  salmon: 1,
  tuna: 2,
  deleteEverything: () => set({ }, true), // clears the entire store, actions included
}))

异步

所有函数都支持异步，因为修改 store 并不依赖返回值，而是调用 set，所以是否异步对数据流框架来说都一样。

监听指定变量

还是用英文比较表意，即 subscribeWithSelector，这个中间件可以让我们把 selector 用在 subscribe 函数上，相比于 redux 传统的 subscribe，就可以有针对性的监听了：

import { subscribeWithSelector } from 'zustand/middleware'
const useStore = create(subscribeWithSelector(() => ({ paw: true, snout: true, fur: true })))

// Listening to selected changes, in this case when "paw" changes
const unsub2 = useStore.subscribe(state => state.paw, console.log)
// Subscribe also exposes the previous value
const unsub3 = useStore.subscribe(state => state.paw, (paw, previousPaw) => console.log(paw, previousPaw))
// Subscribe also supports an optional equality function
const unsub4 = useStore.subscribe(state => [state.paw, state.fur], console.log, { equalityFn: shallow })
// Subscribe and fire immediately
const unsub5 = useStore.subscribe(state => state.paw, console.log, { fireImmediately: true })

后面还有一些结合中间件、immer、localstorage、redux like、devtools、combime store 就不细说了，都是一些细节场景。值得一提的是，所有特性都是正交的。

精读

其实大部分使用特性都在利用 React 语法，所以可以说 50% 的特性属于 React 通用特性，只是写在了 zustand 文档里，看上去像是 zustand 的特性，所以这个库真的挺会借力的。

创建 store 实例

任何数据流管理工具，都有一个最核心的 store 实例。对 zustand 来说，便是定义在 vanilla.ts 文件的 createStore 了。

createStore 返回一个类似 redux store 的数据管理实例，拥有四个非常常见的 API：

export type StoreApi<T extends State> = {
  setState: SetState<T>
  getState: GetState<T>
  subscribe: Subscribe<T>
  destroy: Destroy
}

首先 getState 的实现：

const getState: GetState<TState> = () => state

就是这么简单粗暴。再看 state，就是一个普通对象：

let state: TState

这就是数据流简单的一面，没有魔法，数据存储用一个普通对象，仅此而已。

接着看 setState，它做了两件事，修改 state 并执行 listenser：

const setState: SetState<TState> = (partial, replace) => {
  const nextState = typeof partial === 'function' ? partial(state) : partial
  if (nextState !== state) {
    const previousState = state
    state = replace ? (nextState as TState) : Object.assign({}, state, nextState)
    listeners.forEach((listener) => listener(state, previousState))
  }
}

修改 state 也非常简单，唯一重要的是 listener(state, previousState)，那么这些 listeners 是什么时候注册和声明的呢？其实 listeners 就是一个 Set 对象：

const listeners: Set<StateListener<TState>> = new Set()

注册和销毁时机分别是 subscribe 与 destroy 函数调用时，这个实现很简单、高效。对应代码就不贴了，很显然，subscribe 时注册的监听函数会作为 listener 添加到 listeners 队列中，当发生 setState 时便会被调用。

最后我们看 createStore 的定义与结尾：

function createStore(createState) {
  let state: TState
  const setState = /** ... */
  const getState = /** ... */
  /** ... */
  const api = { setState, getState, subscribe, destroy }
  state = createState(setState, getState, api)
  return api
}

虽然这个 state 是个简单的对象，但回顾使用文档，我们可以在 create 创建 store 利用 callback 对 state 赋值，那个时候的 set、get、api 就是上面代码倒数第二行传入的：

import { create } from 'zustand'

const useStore = create((set, get) => ({
  bears: 0,
  increasePopulation: () => set(state => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 })
}))

至此，初始化 store 的所有 API 的来龙去脉就梳理清楚了，逻辑简单清晰。

create 函数的实现

上面我们说清楚了如何创建 store 实例，但这个实例是底层 API，使用文档介绍的 create 函数在 react.ts 文件定义，并调用了 createStore 创建框架无关数据流。之所 create 定义在 react.ts，是因为返回的 useStore 是一个 Hooks，所以本身具有 React 环境特性，因此得名。

该函数第一行就调用 createStore 创建基础 store，因为对框架来说是内部 API，所以命名也叫 api：

const api: CustomStoreApi = typeof createState === 'function' ? createStore(createState) : createState

const useStore: any = <StateSlice>(
  selector: StateSelector<TState, StateSlice> = api.getState as any,
  equalityFn: EqualityChecker<StateSlice> = Object.is
) => /** ... */

接下来所有代码都在创建 useStore 这个函数，我们看下其内部实现：

简单来说就是利用 subscribe 监听变化，并在需要的时候强制刷新当前组件，并传入最新的 state 给到 useStore。所以第一步当然是创建 forceUpdate 函数:

const [, forceUpdate] = useReducer((c) => c + 1, 0) as [never, () => void]

然后通过调用 API 拿到 state 并传给 selector，并调用 equalityFn（这个函数可以被定制）判断状态是否发生了变化：

const state = api.getState()
newStateSlice = selector(state)
hasNewStateSlice = !equalityFn(
  currentSliceRef.current as StateSlice,
  newStateSlice
)

如果状态变化了，就更新 currentSliceRef.current：

useIsomorphicLayoutEffect(() => {
  if (hasNewStateSlice) {
    currentSliceRef.current = newStateSlice as StateSlice
  }
  stateRef.current = state
  selectorRef.current = selector
  equalityFnRef.current = equalityFn
  erroredRef.current = false
})

useIsomorphicLayoutEffect 是同构框架常用 API 套路，在前端环境是 useLayoutEffect，在 node 环境是 useEffect：

说明一下 currentSliceRef 与 newStateSlice 的功能。我们看 useStore 最后的返回值：

const sliceToReturn = hasNewStateSlice
  ? (newStateSlice as StateSlice)
  : currentSliceRef.current
useDebugValue(sliceToReturn)
return sliceToReturn

发现逻辑是这样的：如果 state 变化了，则返回新的 state，否则返回旧的，这样可以保证 compare 函数判断相等时，返回对象的引用完全相同，这个是不可变数据的核心实现。另外我们也可以学习到阅读源码的技巧，即要经常跳读。

那么如何在 selector 变化时更新 store 呢？中间还有一段核心代码，调用了 subscribe，相信你已经猜到了，下面是核心代码片段：

useIsomorphicLayoutEffect(() => {
  const listener = () => {
    try {
      const nextState = api.getState()
      const nextStateSlice = selectorRef.current(nextState)
      if (!equalityFnRef.current(currentSliceRef.current as StateSlice, nextStateSlice)) {
        stateRef.current = nextState
        currentSliceRef.current = nextStateSlice
        forceUpdate()
      }
    } catch (error) {
      erroredRef.current = true
      forceUpdate()
    }
  }
  const unsubscribe = api.subscribe(listener)
  if (api.getState() !== stateBeforeSubscriptionRef.current) {
    listener() // state has changed before subscription
  }
  return unsubscribe
}, [])

这段代码要先从 api.subscribe(listener) 看，这使得任何 setState 都会触发 listener 的执行，而 listener 利用 api.getState() 拿到最新 state，并拿到上一次的 compare 函数 equalityFnRef 执行一下判断值前后是否发生了改变，如果改变则更新 currentSliceRef 并进行一次强制刷新（调用 forceUpdate）。

context 的实现

注意到 context 语法，可以创建多个互不干扰的 store 实例：

import create from 'zustand'
import createContext from 'zustand/context'

const { Provider, useStore } = createContext()

const createStore = () => create(...)

const App = () => (
  <Provider createStore={createStore}>
    ...
  </Provider>
)

首先我们知道 create 创建的 store 是实例间互不干扰的，问题是 create 返回的 useStore 只有一个实例，也没有 <Provider> 声明作用域，那么如何构造上面的 API 呢？

首先 Provider 存储了 create 返回的 useStore：

const storeRef = useRef<TUseBoundStore>()
storeRef.current = createStore()

那么 useStore 本身其实并不实现数据流功能，而是将 <Provider> 提供的 storeRef 拿到并返回：

const useStore: UseContextStore<TState> = <StateSlice>(
  selector?: StateSelector<TState, StateSlice>,
  equalityFn = Object.is
) => {
  const useProviderStore = useContext(ZustandContext)
  return useProviderStore(
    selector as StateSelector<TState, StateSlice>,
    equalityFn
  )
}

所以核心逻辑还是是现在 create 函数里，context.ts 只是利用 ReactContext 将 useStore “注入” 到组件，且利用 ReactContext 特性，这个注入可以存在多个实例，且不会相互影响。

中间件

中间件其实不需要怎么实现。比如看这个 redux 中间件的例子：

import { redux } from 'zustand/middleware'
const useStore = create(redux(reducer, initialState))

可以将 zustand 用法改变为 reducer，实际上是利用了函数式理念，redux 函数本身可以拿到 set, get, api，如果想保持 API 不变，则原样返回 callback 就行了，如果想改变用法，则返回特定的结构，就是这么简单。

为了加深理解，我们看看 redux 中间件源码：

export const redux = ( reducer, initial ) => ( set, get, api ) => {
  api.dispatch = action => {
    set(state => reducer(state, action), false, action)
    return action
  }
  api.dispatchFromDevtools = true
  return { dispatch: (...a) => api.dispatch(...a), ...initial }
}

将 set, get, api 封装为 redux API：dispatch 本质就是调用 set。

总结

zustand 是一个实现精巧的 React 数据流管理工具，自身框架无关的分层合理，中间件实现巧妙，值得学习。