项目中使用了zustand，那它的原理是什么？

真正的粉丝
66 阅读5分钟

zustand 作为现今使用最为广泛的状态管理工具，本文从 zustand 源码实现解析原理，考虑并非所有读者都了解状态管理工具的实现方法，本文会讲解得讲解的比较详细

zustand 代码案例

先看一段 zustand 使用案例

// store.ts
import { create } from 'zustand';
interface DemoStore {
  count: number;
  setCount: (count: number) => void;
}
export const useDemoStore = create<DemoStore>((set, get) => ({
  count: 1,
  setCount: (count: number) => set({ count }),
}));

// 业务代码中
const [count, setCount] = useDemoStore(state => [state.count, state.setCount]);

简单来说就是个 useDemoStore，里面包含 count 变量和变量变更的方法

初探 create

首先可以看出来的的是，业务代码中从 zustand 导入的实际使用的接口是 create，那么在 zustand 源码 react.ts 中可以找到 create 方法，简化后为

export const create = (createState) => createState
    ? createImpl(createState)
    : createImpl

这种写法是经典的函数重载 + 柯里化兼容，在库设计中很常见，目的在于既支持直接调用，又支持类型安全的延迟调用；

进一步简化，可以初略认为 create = createImpl

const createImpl = createState => {
  const api = createStore(createState)
  const useBoundStore = selector => useStore(api, selector)
  Object.assign(useBoundStore, api)
  return useBoundStore
}

createImpl 中有两步关键步骤

  1. 一个是通过 createStore 创造一个 api，这个 api 的类型是 StoreApi；可以把 zustand store 想象成一个小型数据库或服务，用户不需要知道 store 内部的实现，只要通过 api 进行操控就行

  2. 另一个是 useStore 实现 react 渲染

createStore - 怎么做到控制状态变更的

要看如何做到控制状态更新，要先说 createStore，其简化代码如下

const createStore = (createState) => {
  let state;
  const listeners = new Set();

  const setState = (partial, replace) => {
    // 计算下一个状态
    const nextState = typeof partial === 'function'
      ? partial(state)  // 函数式更新：fn(state)
      : partial;        // 直接传入新值

    // 如果新旧状态相同（用 Object.is 比较），不更新
    if (!Object.is(nextState, state)) {
      const previousState = state;

      // 决定是否完全替换状态
       state =
        (replace ??
         (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)

      // 通知所有订阅者：状态变了
      listeners.forEach(listener => listener(state, previousState));
    }
  };

  const getState = () => state;

  const getInitialState = () => initialState;

  const subscribe = (listener) => {
    listeners.add(listener);
    // 返回取消订阅函数
    return () => listeners.delete(listener);
  };

  // 创建 API 对象
  const api = { setState, getState, getInitialState, subscribe };

  // 初始化状态：调用用户传入的 createState 函数
  const initialState = (state = createState(setState, getState, api));

  // 返回 store 的 API
  return api;
};

createState 作为形参，从 create 一直传入 createStore，其从实质上来说就是我们业务代码在设置 store.ts 时的配置

(set, get) => ({
  count: 1,
  setCount: (count: number) => set({ count }),
})

可以注意到在 createStore 结束前，会赋值 initialState 和 state 并返回 api 供外部操作

const initialState = (state = createState(setState, getState, api))
return api as any

也就是说引用zustand的代码中的 set 方法会被对应到 StoreApi 的 setState 方法， get 方法会对应到 StoreApi 的 getState 方法，在 store.ts 中通过 set 方法变更状态时会调用 setState 使闭包中的 state 变更，并通知闭包中的 listeners 也就是订阅者

setState - 变更完成后，如何判断是否更新

已经知道了是怎么做到操作 store 的，那么 zustand 是判断 state 是否更新的呢？又是怎么避免不必要的更新的？

可以查看 setState 方法（简化方法，不考虑 replace）

const setState = (partial) => {
  const nextState = typeof partial === 'function'
    ? partial(state)
    : partial
  // 如果新状态和旧状态相同，不更新（避免无意义触发）
  if (!Object.is(nextState, state)) {
    const previousState = state
    state = (typeof nextState !== 'object' || nextState === null)
      ? nextState
      : Object.assign({}, state, nextState) // 兼容性考虑

    // 通知所有订阅者：状态已更新
    listeners.forEach((listener) => listener(state, previousState))
  }
}

可以看出，每次 setState 执行完成后会比对新旧 state 是否变化，如果不同会通知所有的订阅者更新，如果相同就不会进行更新，这样避免了不必要的更新

useStore - state 变更完成后，如何通知 react 进行渲染

这就要说到另一个核心了，就是 useStore，通过 createStore 获取的 StoreApi 会通过 useStore 保证变更时渲染

export function useStore(api, selector) {
  const slice = React.useSyncExternalStore(
    api.subscribe, // 订阅
    () => selector(api.getState()), // 当前组件中需要的 store 数据快照
    () => selector(api.getInitialState()), // SSR时的初始值
  )
  React.useDebugValue(slice)
  return slice
}

useStore 的实质是 useSyncExternalStore，这是 react18 后新hook的，作用在于方便订阅外部 store

每次在页面中引入 useDemoStore 时都是一次订阅，如果 setState 判断 state 发生变更，那么根据 setState 结束时的操作会通知所有订阅者（就是listener）获取当前快照，然后按照快照是否变化判断渲染页面

() => selector(api.getState()) // 当前组件中需要的 store 数据快照

其中 selector 是每个业务代码中引入 useStore 中申明的，比如 useDemoStore 中

// 业务代码
const [count, setCount] = useDemoStore(state => [state.count, state.setCount]);

// selector
state => [state.count, state.setCount]

selector 的目的在于获取当前业务代码所需最小的状态切片，只有当前业务代码页面需要的部分状态才会导入页面中，而不用在每个页面中导入完整的 store

如果全量导入 store，会怎么样

存在一种不太正规的写法，就是不使用 selector，直接全量获取 store

const {count, setCount} = useDemoStore();

这样做其实是全量获取 store 然后解构对象，如果 store 变更，这个业务代码页面必定更新

// 全量获取 store 的本质
const state = useDemoState();
const {count, setCount} = state;

by the way，这样做代码能跑通的原因，是 zustand 为了方便我们做了 ts 函数重载处理

const identity = <T>(arg: T): T => arg
export function useStore<S extends ReadonlyStoreApi<unknown>>(
  api: S,
): ExtractState<S>

export function useStore<S extends ReadonlyStoreApi<unknown>, U>(
  api: S,
  selector: (state: ExtractState<S>) => U,
): U

export type ExtractState<S> = S extends { getState: () => infer T } ? T : never

在没有 selector 时，返回的是 api.getState()

总结，到底是什么保证 zustand 没有过度渲染

  1. storeApi 中 setState 会在每次调用时比较新旧 state，只有发生变更才会通知订阅者
  2. 每个页面中引入 store 时，会使用 selector 进行状态切片，保证需要的 state 发生变更才会重新渲染
avatar
文章
阅读
粉丝