一文弄懂Zustand源码实现

前言

最近在为公司项目做技术选型的时候，尝试用了许多款React生态里的状态管理工具，到最后来来去去还是觉得Zustand的轻量和使用方式更加舒服，同时为了能够更好的使用和定位问题，也去尝试阅读了一下源码，结果发现Zustand的源码其实也非常的干净和简洁，特此来利用一篇文章进行一个小小的总结。

前置知识

以下知识点可以帮助你更好的去阅读和理解本篇文章。

React Fiber & Concurrent Mode

在React16之前，React使用的是一种Stack Reconciler。React会在内部维护一个Virtual DOM tree，每当有状态触发更新时，会构建出一棵新的Virtual DOM tree，然后对两棵树进行Diff算法，找到需要修改的地方，修改完毕后一次性更新到真实的Dom上，整个过程是同步且阻塞的。

而在16版本之后，React对jsx代码转换到真实Dom的整个流程进行了优化，将整个更新过程划分为一个个小的工作单元，这个工作单元保存了每个VDom的父节点、子节点和兄弟节点信息，并将整个构建Dom tree和Diff tree的方式从递归变成了循环，这样一来，就可以在每次循环结束后，通过询问浏览器是否存在空余时间的方式，可以选择交出线程控制权或者恢复上一次的Diff工作。

这个小小的工作单元就是Fiber，而这整个工作流程的改动有效的解决了Dom Tree过大的更新卡顿或延迟问题。

那么什么是Concurrent Mode呢？在React引入Fiber架构后，便逐步启用了一个全新的渲染模式，称为并发模式，主要为了解决任务拆分调度和优先级问题。

让我们来思考一个模糊搜索的场景，当页面需要根据用户的输入来实时更新结果列表时，我们需要维护一个Input的Value的状态，以及一个结果列表List的状态。而这两个状态在大部分情况下，都是一起更新的（setState）。如果是在同步模式下，每当触发一个setState，React都会进行一次上述图中的渲染和计算过程，如果有多个状态同时进行setState，则按照触发的顺序，依次进行。但是倘若列表的数据非常大，每次处理的结果都比较慢时，用户就会感觉到十分的卡顿，输入的内容也不能实时的展现出来。

怎么解决呢？方法有很多，其中一种就是指定某个setState的优先级。例如，在这个例子中，我们认为响应用户输入的更新重要程度更高，如果这个state发生变化时，其他的setState需要为这个让路，等待他更新完毕后再进行。而这，实际上就是Concurrent Mode做的事情，也就是基于优先级的可打断的渲染流程。

useSyncExternalStore & useSyncExternalStoreWithSelector

在React中，我们所说的状态通常分为三种：

• 组件内部的State/Props
• 上下文Context
• 组件外部的独立状态Store(Redux/Zustand)

前两种状态实际上都是React内部维护的Api，自然也会跟随着React版本的迭代而进行相对应的优化。 但是组件外部的状态，对于React来说并不可控，如果需要更好的契合React本身，我们需要去写一些与本身业务逻辑无关的胶水代码。

例如：

• 订阅外部状态
• 外部状态更新时，对组件进行重渲染。

export const useOutSideStore = (store) => {
  const [state, setState] = useState(store.getState());

  useEffect(() => {
    const unsubscribe = store.subscribe(() => {
      setState(store.getState());
    });

    return () => { unsubscribe(); };
  }, []);

  return state;
}

但是在React18推出Concurrent Mode后，这种外部状态的订阅模式会存在一个问题，也就是被称为撕裂问题的Bug。

我们来思考这样一种场景： 假设我们的现在的页面触发了更新，需要进行Rerender，而根据Concurrent Mode并发模式下的规则，我们会把更新过程中需要执行的任务划分优先级，优先级低的有可能会被打断。假设某个任务A和B，都同时依赖了外部状态中的某个State，在Rerender开始时，值为1，任务A执行完之后，React把线程的处理权交还给了浏览器，浏览器的某些操作导致了这个State的值变成了2。那么等到B任务重新恢复执行时，读到的值就会出现差异，导致渲染结果的不一致性。

针对上述的一些外部状态与React本身不契合的情况，React提供了一个名为useSyncExternalStore的Hook，这个hook可以让我们更加方便的去订阅外部的Store，并且避免发生撕裂问题。

这个Hook接收三个参数：

• 一个订阅函数，并且该函数返回一个取消订阅的方法。订阅函数的入参实际上就是一个通知组件更新的方法。
• 从store中读取数据的快照，也就是获取Store里的最新数据。

export function useSyncExternalStore<Snapshot>(
  subscribe: (onStoreChange: () => void) => () => void,
  getSnapshot: () => Snapshot,
  getServerSnapshot?: () => Snapshot,
): Snapshot;

而useSyncExternalStoreWithSelector的优化主要是允许从一个大store中取出组件所用到的部分，同时借助isEqual来减少re-render的次数。这个方法维护在了use-sync-external-store/shim/with-selector中。

export function useSyncExternalStoreWithSelector<Snapshot, Selection>(  
  subscribe: (() => void) => () => void,  
  getSnapshot: () => Snapshot,  
  getServerSnapshot: void | null | (() => Snapshot),  
  selector: (snapshot: Snapshot) => Selection,  
  isEqual?: (a: Selection, b: Selection) => boolean,  
): Selection

基本使用

在学习完一些前置知识之后，我们先来简单看下Zustand的基本使用。首先，zustand是一个基于发布订阅模式实现的状态管理方案。我们可以通过create方法去创建一个Store，并在Store里定义我们需要维护的状态和改变状态的方法。

import { create } from "zustand";  
  
const initStateCreateFunc = (set) => ({  
  bears: 0,  
  increase: (by) => set((state) => ({ bears: state.bears + by })),  
});  
  
const useBearStore = create(initStateCreateFunc);

create函数实际上返回了一个Hook，通过调用这个hook，我们就可以在组件中去订阅某个状态，或者获取改变某个状态的方法。

function BearCounter() {  
  const bears = useBearStore((state) => state.bears);  
  return <h1>{bears} around here...</h1>;  
}  
  
function Controls() {  
  const increase = useBearStore((state) => state.increase);  
  return <button onClick={increase}>one up</button>;  
}

源码解析

在Zustand中，基于发布订阅模式的Store创建方法定义在了vanilla.ts中，这个方法并没有耦合任何第三方框架，可以在任意场景下使用。而对于React中的集成，则是放在了react.ts中

createStoreImpl & createStore

我们首先来逐行阅读一下vanilla.ts的核心代码：

// 首先接受一个创建Store的方法，这个方法由使用者定义。
const createStoreImpl: CreateStoreImpl = (createState) => {
  type TState = ReturnType<typeof createState>
  type Listener = (state: TState, prevState: TState) => void
  let state: TState
  // 创建一个Set结构来维护订阅者。
  const listeners: Set<Listener> = new Set()

  // 定义更新数据的方法，partial参数支持对象和函数，replace指的是全量替换store还是merge
  // 如果是partial对象时，则直接赋值，否则将上一次的数据作为参数执行该方法。
  // 然后利用Object.is进行新老数据的浅比较，如果前后发生了改变，则进行替换
  // 并且遍历订阅者，逐一进行更新。
  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        replace ?? typeof nextState !== 'object'
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  // getState方法实则是返回当前Store里的最新数据
  const getState: StoreApi<TState>['getState'] = () => state

  // 添加订阅方法，并且返回一个取消订阅的方法。还记得前置知识里所提到的useSyncExternalStore吗？
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }


  const api = { setState, getState, subscribe }
  // 这里就是官方示例里的set,get,api
  state = createState(setState, getState, api)
  return api as any
}

// 根据传入的createState的类型，手动创建核心store。
export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

接着我们再来看一下在react.ts中，zustand是如何处理的，我们先来看下入口函数（我会把一些干无关的代码也刨除）：

当我们调用create方法时，实际上会先调用上文提到的createStore的方法生成store，然后再利用 useSyncExternalStoreWithSelector方法对react进行集成。

// 对React进行集成
export function useStore<TState, StateSlice>(
  api: WithReact<StoreApi<TState>>,
  selector: (state: TState) => StateSlice = api.getState as any,
  equalityFn?: (a: StateSlice, b: StateSlice) => boolean
) {

  // 利用useSyncExternalStoreWithSelector，对store里的所有数据进行选择性的分片
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getState,
    selector,
    equalityFn
  )
  
  useDebugValue(slice)
  return slice
}

import { createStore } from './vanilla.ts'

const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  // 创建store
  const api =
    typeof createState === 'function' ? createStore(createState) : createState

  // 将store和react进行集成
  const useBoundStore: any = (selector?: any, equalityFn?: any) =>
    useStore(api, selector, equalityFn)

  Object.assign(useBoundStore, api)

  return useBoundStore
}


export const create = (<T>(createState: StateCreator<T, [], []> | undefined) =>
  createState ? createImpl(createState) : createImpl) as Create

总结

通过上述的源码分析，我们可以看到zustand的核心代码其实不多，主要还是基于发布订阅模式去实现了基本的数据维护和通知更新，而与框架的集成也是利用了React官方提供的useSyncExternalStoreWithSelector来实现。