Zustand底层原理解析

Posted on Edited on

介绍

Zustand是最近比较火的状态管理库，是Redux的替代品，相比于Redux，它的api简单，样板代码少，包体积小只有1.2KB（minified），基于hooks来管理状态以及生态也比较丰富，有immer、persist、redux等中间件。

原理

src/vanilla.ts中的代码如下，vanilla的意思是这里导出的createStore各框架都能用，而非仅限于React。
源码比较简单，核心是基于发布订阅模式和useSyncExternalStore实现的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
// src/vanilla.ts
const createStoreImpl: CreateStoreImpl = (createState) => {
  type TState = ReturnType<typeof createState>
  type Listener = (state: TState, prevState: TState) => void
  let state: TState
  // 订阅者，包含了每个组件通过useSyncExternalStore绑定的重新渲染的方法
  const listeners: Set<Listener> = new Set()

  // 变更状态，可以只变更部分state，会做合并
  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 === null))
          ? (nextState as TState)
          // 更新state时用的是Object.assign，更改了state的引用，所以在使用state时，需要用selector或者useShallow。
          : Object.assign({}, state, nextState)
      // setState时触发更新
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

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

  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
    initialState

  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }

  const api = { setState, getState, getInitialState, subscribe }
  const initialState = (state = createState(setState, getState, api))
  return api as any
}

export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

在React中集成

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
// src/react.ts
export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  // 挂载组件更新的方法
  const slice = React.useSyncExternalStore(
    api.subscribe,
    () => selector(api.getState()),
    () => selector(api.getInitialState()), // getServerSnapshot
  )
  React.useDebugValue(slice)
  return slice
}

const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState)

  const useBoundStore: any = (selector?: any) => useStore(api, selector)

  // 在store上挂载setState, getState, getInitialState, subscribe 这些方法
  Object.assign(useBoundStore, api)

  return useBoundStore
}

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

可以看到，Zustand是对useSyncExternalStore进行了封装，将状态统一管理在组件外部，避免了层层传递 Props。这种方式使得状态管理更加灵活，开发过程更加高效。

与Redux的对比

  1. Zustand相对于Redux需要较少的代码，并且上手成本低
  2. 定义衍生状态，Redux通过自定义hook/或者mapStateToProps实现，Zustand可以定义computed属性，相对来说比较方便。
  3. Redux是单一store，Zustand可以是多Store

与MobX的对比

  1. 没有computed缓存
  2. 更新粒度没有MobX细，开发时需要考虑如何避免重复渲染。
  3. 代码相对于MobX Store来说不够直观
  4. 体积小，Zustand + immer之后是14KB，mobx + mobx-react + mobx-react-lite是35KB
  5. 设置新状态需要用set(state => newState)的方式，没有MobX来的直接并且符合直觉。
  6. 避免re-render不需要selector，心智负担更小，代码量更少。

如果不考虑体积，我觉得可以直接选MobX。