zustand 源码学习前言 zustand 源代码出乎意料的精简，适合学习阅读，大家可以下载源码并自行阅读。源码地址

前言

zustand 源代码出乎意料的精简，适合学习阅读，大家可以下载源码并自行阅读。

入口

直接看 index.ts 源码：

export * from './vanilla.ts'
export * from './react.ts'

可以看到入口文件主要起了一个集中导出的作用，将 vanilla.ts 和 react.ts 中的资源一起导出。

`vanilla.ts` 和 `react.ts`

先说结论，vanilla.ts 是基础，react.ts 是在 vanilla.ts 基础上做的 react 支持。

`vanilla.ts`

先贴代码：


const createStoreImpl: CreateStoreImpl = (createState) => {
  let state: TState
  const listeners: Set<Listener> = new Set()

  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)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  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 框架相关知识，这段代码很容易理解。

createStore 是一个接收 createState 参数的函数，这个函数根据是否传入 createState 参数返回不同的结果。如果存在 createState ，则执行 createStoreImpl 函数并传入 createState 如果没有传入，则直接返回 createStoreImpl 函数。

那重点就是 createStoreImpl 函数，在这个函数中定义了 4 个函数。分别是：

setState: 设置状态
getState: 获取状态
getInitialState: 获取初始状态
subscribe: 订阅

这 4 个函数从函数名可以很轻松的知道它们的作用。

那这里的状态指的是什么呢？

当然指的是 createStoreImpl 函数定义的内部变量，state。并且通过执行 createState 来初始化这个状态，而 createState 的参数就是上面定义的 setState, getState, api。api 为一个对象，包含了这4 个函数。最后返回 api。

订阅功能维护了一个数组 listeners，用来保存所有订阅者，在 set 时触发所有订阅。subscribe 函数会返回一个取消订阅函数。

调用实例

const useStoreA = createStore<Store>((set, get) => ({
  count: 0,
  inc: () => set({ count: get().count + 1 }),
}))

那么到这里就会产生一个疑问，如何处罚 react 更新呢？答案就在 react.ts 文件中。

`react.ts`

同样直接看代码：

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

const identity = <T>(arg: T): T => arg

export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  const slice = React.useSyncExternalStore(
    api.subscribe,
    React.useCallback(() => selector(api.getState()), [api, selector]),
    React.useCallback(() => selector(api.getInitialState()), [api, selector]),
  )
  React.useDebugValue(slice)
  return slice
}

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

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

  Object.assign(useBoundStore, api)

  return useBoundStore
}

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

这个文件一样非常简单，它导出了两个函数：

useStore
create

其中 create 函数和 createStore 函数的作用基本一致，这里不细说。直接看 createImpl 函数。

这个函数就 4 行代码，非常好理解。

调用 vanilla.ts 中的 createStore 函数得到 api 对象。
定义 useBoundStore 函数
合并 api 对象和 useBoundStore 函数
返回 useBoundStore 函数

其中合并 api 对象和 useBoundStore 函数的操作让我们可以将 useBoundStore 当做 hook 使用，也可以在组件外部当 store 使用，如 useBear.getState().bears。

所以，在 useBoundStore 函数中调用的 useStore 函数是关键。

useStore 函数主要就是调用了 React.useSyncExternalStore 并返回结果，其中使用 React.useCallback 做了优化，保证依赖稳定，并且使用选择器监听指定元素。如：

const count = useStore((s) => s.count)

到这里 zustand 的主逻辑其实就基本结束了，但是会有一个问题，如何扩展？

那么就需要了解 zustand 的插件系统。

插件

要了解 zustand 的插件系统设计，首先需要知道为什么会设计插件？

毫无疑问是需要对 zustand 功能在核心功能上进行增强。比如在写入之前做日志，批量更新，静态化，增加新方法等等。

zustand 采用了中间件的形式来实现。

对于中间件有几种模式：

装饰器模式

// (config, options?) => (set, get, api) => state
const persist = (config, opts) => (set, get, api) => {
  const wrappedSet = (...args) => { /* 持久化 */; return set(...args) }
  return config(wrappedSet, get, api)
}

createStore(persist(createState, { name: 'app' }))

工厂式

// 例：redux(reducer, initialState) => StateCreator
const store = createStore(redux(reducer, initialState))

柯里化式

const logger = (opts) => (config) => (set, get, api) => {
  const wrappedSet = (...args) => { /* 打日志 */; return set(...args) }
  return config(wrappedSet, get, api)
}

createStore(logger({ name: 'counter' })(createState))

创建后增强

在createStore 之后直接改造 api（不包 stateCreator）：

const store = createStore(createState)
const enhance = (api) => {
  const orig = api.setState
  api.setState = (patch, replace) => { /* 统计/埋点 */; return orig(patch, replace) }
}
enhance(store)

zustand 源码学习

前言

入口

vanilla.ts 和 react.ts

vanilla.ts