深入Zustand系列- Zustand源码解析

上篇文章用万字长文给大家带来为什么我们需要Zustand，今天我们换个视角，一起来看看Zustand内部是如何实现的。

Zustand内部使用了分层架构和中间件模式，将核心逻辑和扩展功能分开，实现了模块化和高可扩展性。主要由三个核心文件组成:

vanilla.ts：核心store的实现
react.ts：连接react环境
Middleware: 中间件

核心状态管理逻辑实现

和redux类似，Zustand 状态管理的内部核心逻辑是一个与UI框架解耦、可以独立运行的Javascript函数。这个函数简单到你甚至可以在node.js上运行，而所有的实现都放在 vanilla.ts

vanilla.ts

typescript

const createStoreImpl: CreateStoreImpl = (createState) => {
  // 1. 类型推断：从 createState 函数推断状态类型
  type TState = ReturnType<typeof createState>
  
  // 2. 定义监听器类型：接收新状态和旧状态
  type Listener = (state: TState, prevState: TState) => void
  
  // 3. 状态变量：存储当前状态
  let state: TState
  
  // 4. 监听器集合：使用 Set 存储所有订阅者
  const listeners: Set<Listener> = new Set()

  // 5. setState 函数实现
  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    // 5.1 计算新状态
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)  // 函数式更新
        : partial  // 直接对象更新
    
    // 5.2 性能优化：只有状态真正变化时才更新
    if (!Object.is(nextState, state)) {
      const previousState = state
      
      // 5.3 状态更新逻辑
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)  // 完全替换
          : Object.assign({}, state, nextState)  // 浅合并
      
      // 5.4 通知所有监听器
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  // 6. getState 函数：返回当前状态
  const getState: StoreApi<TState>['getState'] = () => state

  // 7. getInitialState 函数：返回初始状态
  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
    initialState

  // 8. subscribe 函数：订阅状态变化
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // 返回取消订阅函数，利用闭包特性
    return () => listeners.delete(listener)
  }

  // 9. 创建 API 对象
  const api = { setState, getState, getInitialState, subscribe }
  
  // 10. 关键步骤：初始化状态
  const initialState = (state = createState(setState, getState, api))
  
  // 11. 返回 API 对象
  return api as any
}

在状态管理中，通过一个闭包 let state来实现内部的状态存储，使用 Set 来存储 listeners，解决了重复订阅的问题，同时Object.assign({}, state, nextState) 创建一个新的对象确保了state的浅层不可变。

需要非常注意的是，Zustand内部 const initialState = (state = createState(setState, getState, api)) 代码实现了控制反转的设计模式即用户决定状态结构，Store提供了基础设施。Zustand在业务代码中注入了三个关键的内部api

setState: 状态更新函数
getState: 状态获取函数
api: 完整的Store API

这让业务可以通过下面的方式去做状复杂态管理:

javascript

const useStore = create((set, get, api) => ({
  users: [],
  loading: false,
  fetchUsers: async () => {
    set({ loading: true })
    const users = await fetch('/api/users').then(r => r.json())
    set({ users, loading: false })
  }
}))

react.ts

react.ts 是连接 Vanilla 核心和 React 组件的桥梁。下面是剥离typescript之后的react.ts源代码，有一说一，Zustand内部typescript代码写的非常的精妙，会专门写篇文档来记录这些。

javascript

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

const identity = (arg) => arg

export function useStore(api, selector = identity) {
  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 = (createState) => {
  const api = createStore(createState)

  const useBoundStore = (selector) => useStore(api, selector)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

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

create(createState) 本质上是一个包装器，其内部流程如下:

调用vanilla.ts 中的 createStore(createState) 创建一个初始化的store实例
将store实例传递给useStore，并返回一个hook
将store实例上的所有方法（getState, setState, subscribe）直接附加到返回的 useStore Hook 上，就可以在脱离react环境直接使用，实现了"既是Hook又是Store API"的双重身份。比如我们可以通过useCount.getState()也可以在 React hook中useCount()使用;

useStore Hook (响应式核心):

Zustand 的响应式并非基于 Proxy ，也不是基于不可变数据的diff。它是基于发布-订阅模式的主动通知。当 setState 被调用时，它会遍历所有通过 subscribe 注册的监听器（在 React 中就是组件的更新函数），并执行它们。组件是否真的重渲染，则取决于 React 的 useSyncExternalStore Hook 和你提供的 selector 函数。只有当 selector 的返回值发生变化时，组件才会更新，从而实现了精细的渲染。

从Zustand V4开始，useStore 的实现完全依赖React18的useSyncExternalStore,而这也是 React 官方推荐的、用于订阅外部数据源（如状态管理库）的标准方式。大家可在这里看 useSyncExternalStore 实操指南。

javascript

export function useStore(api, selector = identity) {
  const slice = React.useSyncExternalStore(
    // 1. 订阅函数：当外部状态变化时调用
    api.subscribe,
    // 2. 获取当前状态的函数
    React.useCallback(() => selector(api.getState()), [api, selector]),
    // 3. 获取初始状态的函数（SSR支持）
    React.useCallback(() => selector(api.getInitialState()), [api, selector]),
  )
  React.useDebugValue(slice)
  return slice
}

api.subscribe: useSyncExternalStore 会调用这个函数来注册一个当 store 发生变化时需要执行的回调。Zustand 内部的 setState 在状态更新后会触发这个回调，通知 React 可能需要进行重渲染。
api.getState: 用于在任何时间点获取 store 的最新状态快照。React 会在初始渲染和每次更新时调用它来获取当前状态。
selector 函数: useSyncExternalStore 获取到完整的 state 后，会立即用 selector 对其进行处理（例如 state => state.bears）。只有 selector 的返回值会从 Hook 中返回
equalityFn 比较函数: React 会用这个函数（默认为 Object.is）来比较 selector 在上一次渲染和这一次渲染中的返回值。如果返回值没有变化（equalityFn 返回 true），React 就会跳过这次重渲染。这就是 Zustand 高性能的基石：它确保了只有在组件真正依赖的状态切片发生变化时，组件才会更新。

middleware

在日常中，我们可以通过下面的使用姿势去使用devtools的中间件。

javascript

import { create } from  'zustand'
import { devtools } from 'zustand/middleware'

// 基础使用
const useStore = create(
  devtools(
    (set, get) => ({
      count: 0,
      increment: () => set((state) => ({ count: state.count + 1 }))
    }),
    { name: 'counter-store' } // DevTools中显示的名称
  )
)

// 在组件中使用
function Counter() {
  const { count, increment } = useStore()
  
  return (
    <div>
      <span>{count}</span>
      <button onClick={increment}>+1</button>
    </div>
  )
}

在核心实现上，Zustand 采用了洋葱模式实现中间件的架构，通过高阶函数、类型系统和函数组合搭配实现了完整的功能。具体的架构如下图所示:

Zustand内部提供了immer、devtools、redux、persist等中间件，其内部实现都大同小异。我们可以通过这个公用的套路去实现一个自己的中间件。

如何实现一个中间件

1. 类型定义和模块声明

定义中间件需要的类型、常量、工具类
处理配置选项的合并和验证
初始化中间件内部状态

2. 增强 API（可选）

向 api 对象添加新的方法
扩展 store 的功能接口
为外部提供中间件特定的 API

3. 重写方法（可选）

重写 api.setState、api.getState 等方法
在方法调用前后添加自定义逻辑
实现状态拦截和增强

4. 调用原始函数

调用用户定义的状态创建函数 fn
传入增强后的 set、get、api 参数
获取用户创建的初始状态

5. 返回增强后的状态

返回用户创建的状态对象
可选择性地添加中间件特定的属性
完成中间件的包装过程

通过上面的步骤，我们就可以自己定义一个自己的中间件，下面我们将举一个具体的例子。

下面的例子非常简单，在状态变化时输出日志，显示状态变化前后的state的值和action信息。

javascript


const logImpl = (fn, options = {}) => (set, get, api) => {
  // 1. 类型定义和模块声明
  const config = {
    enabled: true,
    collapsed: true,
    ...options
  }

  // 2. 增强 API（可选）
  // 不需要增强 API

  // 3. 重写方法（可选）
  const originalSetState = api.setState
  api.setState = (state, replace, action) => {
    if (config.enabled) {
      const prevState = get()
      
      if (config.collapsed) {
        console.groupCollapsed(`[Log] ${action || 'anonymous'}`)
      } else {
        console.group(`[Log] ${action || 'anonymous'}`)
      }
      
      console.log('prev state:', prevState)
      console.log('action:', action)
    }

    const result = originalSetState(state, replace, action)

    if (config.enabled) {
      console.log('next state:', get())
      console.groupEnd()
    }

    return result
  }

  // 4. 调用原始函数
  const initialState = fn(api.setState, get, api)

  // 5. 返回增强后的状态
  return initialState
}

export const log = logImpl

而使用log中间件也非常简单，直接如下使用就行。

javascript

import { create } from 'zustand'
import { log } from './middleware/log'

const useStore = create(
  log(
    (set, get) => ({
      // 状态定义
    }),
    {
      enabled: true,    // 是否启用
      collapsed: false  // 是否折叠日志
    }
  )
)

Zustand的源代码通篇读下来，我们就会惊喜的发现。Zustand的优雅是建立在三个务实、简单的模式之上：

用最基础的闭包实现一个全局的store，这解决了存什么的问题
基于发布-订阅的模式结合react的useSyncExternalStore的hook实现了响应式更新核心，这解决了如何更新的问题
通过函数式编程思想（工厂模式、控制反转）巧妙的实现了中间件，这解决了怎么扩展的问题。

熟悉了这些模式，我们也可以实现一个自己的react store库，而这个将在下篇文章中进行详细介绍。

深入Zustand系列- Zustand源码解析 ​

核心状态管理逻辑实现 ​

vanilla.ts ​

react.ts ​

middleware ​

如何实现一个中间件 ​

1. 类型定义和模块声明 ​

2. 增强 API（可选） ​

3. 重写方法（可选） ​

4. 调用原始函数 ​

5. 返回增强后的状态 ​

深入Zustand系列- Zustand源码解析

核心状态管理逻辑实现

vanilla.ts

react.ts

middleware

如何实现一个中间件

1. 类型定义和模块声明

2. 增强 API（可选）

3. 重写方法（可选）

4. 调用原始函数

5. 返回增强后的状态