Zustand 源码学习笔记
本文基于 `src` 目录梳理 Zustand 的核心原理。Zustand 的实现很小，关键思想可以概括为：
1. `vanilla.ts` 提供一个与框架无关的外部 store。
2. `react.ts` 用 `useSyncExternalStore` 把外部 store 接入 React 渲染模型。
3. middleware 通过包裹 `StateCreator`，在初始化阶段扩展或改写 `setState`、`subscribe`、`api`。
4. TypeScript 类型系统用 `StoreMutators` 和 `Mutate` 把 middleware 对 store API 的增强串起来。
目录结构
```text
src/
index.ts # 主入口：导出 vanilla 与 react API
vanilla.ts # 核心 store 实现，不依赖 React
react.ts # React hook 绑定：create/useStore
traditional.ts # 支持 equalityFn 的传统入口
shallow.ts # shallow 与 useShallow 的组合导出
vanilla/shallow.ts # 浅比较算法
react/shallow.ts # 复用 selector 返回值，减少重渲染
middleware.ts # middleware 聚合导出
middleware/*.ts # persist/devtools/redux/immer 等增强
```
一、核心模型：一个外部 Store
Zustand 最核心的文件是 `src/vanilla.ts`。它定义的 store API 只有四个方法：
```ts
export interface StoreApi<T> {
setState: SetStateInternal<T>
getState: () => T
getInitialState: () => T
subscribe: (listener: (state: T, prevState: T) => void) => () => void
}
```
这说明 Zustand 的运行时模型非常直接：
- `state` 是闭包中的一个变量。
- `listeners` 是一个 `Set`，保存所有订阅者。
- `setState` 修改 `state` 后通知订阅者。
- `getState` 读取当前状态。
- `getInitialState` 保存初始化时的状态，主要给 React SSR/hydration 使用。
- `subscribe` 注册监听器，并返回取消订阅函数。
核心实现如下：
```ts
const createStoreImpl: CreateStoreImpl = (createState) => {
type TState = ReturnType<typeof createState>
type Listener = (state: TState, prevState: TState) => void
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)
return () => listeners.delete(listener)
}
const api = { setState, getState, getInitialState, subscribe }
const initialState = (state = createState(setState, getState, api))
return api as any
}
```
这里有几个关键点：
1. `createState` 是 store 的初始化函数
用户写的代码通常是：
```ts
const useStore = create((set, get) => ({
count: 0,
inc: () => set((state) => ({ count: state.count + 1 })),
}))
```
其中 `(set, get) => ({ ... })` 就是源码里的 `createState`，类型名叫 `StateCreator`。
初始化时：
```ts
const initialState = (state = createState(setState, getState, api))
```
这行做了两件事：
1. 调用用户提供的初始化函数，得到初始状态对象。
2. 同时把结果赋给闭包变量 `state` 和 `initialState`。
因此 action 也只是 state 对象上的函数。Zustand 没有 reducer、dispatch 或 action queue 这类强约束模型。
2. `setState` 支持对象和函数两种写法
```ts
const nextState =
typeof partial === 'function'
? partial(state)
: partial
```
所以这两种写法都成立：
```ts
set({ count: 1 })
set((state) => ({ count: state.count + 1 }))
```
3. 默认浅合并，`replace` 为 true 时整体替换
```ts
state =
(replace ?? (typeof nextState !== 'object' || nextState === null))
? nextState
: Object.assign({}, state, nextState)
```
默认情况下，如果 `nextState` 是对象，Zustand 会执行一层浅合并：
```ts
set({ count: 1 })
// 等价于 Object.assign({}, oldState, { count: 1 })
```
如果传 `replace: true`，或 `nextState` 不是对象，则直接替换整个 state：
```ts
setState(newState, true)
```
这也是 Zustand 文档中强调“嵌套对象仍需要不可变更新”的原因：它只帮你合并第一层。
4. 只有状态引用变化时才通知
```ts
if (!Object.is(nextState, state)) {
// update and notify
}
```
如果 `nextState` 和当前 `state` 是同一个引用，监听器不会触发。这是最底层的变化检测。
二、React 绑定：`useSyncExternalStore`
`src/react.ts` 的职责是把 vanilla store 变成 React hook。
核心代码：
```ts
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
}
```
`useSyncExternalStore` 是 React 官方提供的外部 store 订阅 API。三个参数分别是：
- `subscribe`：告诉 React 如何订阅外部 store。
- `getSnapshot`：React 渲染时如何读取当前快照。
- `getServerSnapshot`：SSR/hydration 时如何读取初始快照。
Zustand 的 selector 也是在这里生效的：
```ts
selector(api.getState())
```
如果用户写：
```ts
const count = useStore((state) => state.count)
```
React 组件实际订阅的是整个 store，但渲染时拿到的是 selector 返回的 slice。
三、`create` 为什么既是 Hook 又有 Store API
`react.ts` 里还有一段非常关键：
```ts
const createImpl = <T>(createState: StateCreator<T, [], []>) => {
const api = createStore(createState)
const useBoundStore: any = (selector?: any) => useStore(api, selector)
Object.assign(useBoundStore, api)
return useBoundStore
}
```
这解释了 Zustand 的一个经典用法：
```ts
const useStore = create(...)
useStore((state) => state.count) // React hook
useStore.getState() // 直接读 store
useStore.setState(...) // React 外部更新
useStore.subscribe(...) // 直接订阅
```
原因是 `useBoundStore` 本质上是一个函数，但源码用 `Object.assign` 把 vanilla store 的 API 挂到了这个函数对象上。
所以 `create` 的产物不是单纯的 hook，而是：
```text
function useBoundStore(selector?) {
return useStore(api, selector)
}
useBoundStore.setState = api.setState
useBoundStore.getState = api.getState
useBoundStore.subscribe = api.subscribe
useBoundStore.getInitialState = api.getInitialState
```
四、selector 与 equalityFn
默认入口 `src/react.ts` 使用 React 原生的 `useSyncExternalStore`。它会比较 snapshot 是否变化，但没有额外传入自定义 equality function。
`src/traditional.ts` 提供了 `createWithEqualityFn` 和 `useStoreWithEqualityFn`：
```ts
const slice = useSyncExternalStoreWithSelector(
api.subscribe,
api.getState,
api.getInitialState,
selector,
equalityFn,
)
```
这里使用的是 `use-sync-external-store/shim/with-selector`，多了两个能力：
- selector 由 shim 管理。
- 可以传入 `equalityFn`，决定 slice 是否真的变化。
例如：
```ts
const value = useStoreWithEqualityFn(
store,
(state) => [state.a, state.b],
shallow,
)
```
如果 selector 每次都返回新数组，默认 `Object.is` 会认为它变了；加上 `shallow` 后，数组元素没变就可以跳过重渲染。
五、`shallow` 的浅比较逻辑
`src/vanilla/shallow.ts` 实现了 Zustand 的浅比较。它不只比较普通对象，也处理 `Map`、`Set`、数组等 iterable。
核心流程：
```ts
export function shallow<T>(valueA: T, valueB: T): boolean {
if (Object.is(valueA, valueB)) {
returntrue
}
if (
typeof valueA !== 'object' ||
valueA === null ||
typeof valueB !== 'object' ||
valueB === null
) {
returnfalse
}
if (Object.getPrototypeOf(valueA) !== Object.getPrototypeOf(valueB)) {
returnfalse
}
if (isIterable(valueA) && isIterable(valueB)) {
if (hasIterableEntries(valueA) && hasIterableEntries(valueB)) {
return compareEntries(valueA, valueB)
}
return compareIterables(valueA, valueB)
}
return compareEntries(
{ entries: () => Object.entries(valueA) },
{ entries: () => Object.entries(valueB) },
)
}
```
判断顺序是：
1. `Object.is` 相等，直接返回 `true`。
2. 任意一个不是对象或为 `null`，返回 `false`。
3. 原型不同，返回 `false`，避免把不同类型的对象误判为相同。
4. 如果是 iterable：
- 有 `entries()` 的按 key/value 比较，适合 `Map`、`Set`。
- 没有 `entries()` 的按迭代顺序比较。
5. 普通对象转成 `Object.entries` 后逐项比较。
`src/react/shallow.ts` 的 `useShallow` 不是比较 state 本身，而是包装 selector：
```ts
export function useShallow<S, U>(selector: (state: S) => U): (state: S) => U {
const prev = React.useRef<U>(undefined)
return (state) => {
const next = selector(state)
return shallow(prev.current, next)
? (prev.current as U)
: (prev.current = next)
}
}
```
如果新的 selector 结果和上一次浅相等，它会返回上一次的引用，从而让 React 看到“snapshot 没变”。
六、middleware 的本质
Zustand middleware 的共同模式是：
```ts
middleware(config) => (set, get, api) => {
// 改写 set/get/api
return config(newSet, get, api)
}
```
也就是说 middleware 并不创建另一个 store，它只是包裹用户的 `StateCreator`，在 store 初始化时介入。
### `combine`
`combine` 最简单：
```ts
export function combine(initialState, create) {
return (...args) => Object.assign({}, initialState, create(...args))
}
```
它把初始状态和 create 返回的 action 合并，主要是改善类型推导和组织代码。
### `subscribeWithSelector`
它增强的是 `api.subscribe`：
```ts
api.subscribe = ((selector: any, optListener: any, options: any) => {
let listener: Listener = selector // if no selector
if (optListener) {
const equalityFn = options?.equalityFn || Object.is
let currentSlice = selector(api.getState())
listener = (state) => {
const nextSlice = selector(state)
if (!equalityFn(currentSlice, nextSlice)) {
const previousSlice = currentSlice
optListener((currentSlice = nextSlice), previousSlice)
}
}
if (options?.fireImmediately) {
optListener(currentSlice, currentSlice)
}
}
return origSubscribe(listener)
}) as any
```
增强前只能这样订阅：
```ts
store.subscribe((state, prevState) => {})
```
增强后可以订阅某个 slice：
```ts
store.subscribe(
(state) => state.count,
(count, prevCount) => {},
)
```
内部仍然是订阅整个 store，只是在 listener 里执行 selector，并用 `equalityFn` 判断是否调用用户回调。
### `persist`
`persist` 的核心是持久化和 hydration。
`createJSONStorage` 把普通 string storage 包装成对象 storage：
```ts
const persistStorage = {
getItem: (name) => JSON.parse(storage.getItem(name)),
setItem: (name, newValue) =>
storage.setItem(name, JSON.stringify(newValue)),
removeItem: (name) => storage.removeItem(name),
}
```
`persistImpl` 做了三件关键事：
1. 改写 `api.setState`，每次 set 后写入 storage。
2. 初始化时从 storage 读取旧数据，处理版本迁移。
3. 提供 `store.persist` API，例如 `rehydrate`、`clearStorage`、`hasHydrated`。
关键片段：
```ts
const savedSetState = api.setState
api.setState = (state, replace) => {
savedSetState(state, replace as any)
return setItem()
}
const configResult = config(
(...args) => {
set(...args)
return setItem()
},
get,
api,
)
```
这里同时覆盖了两条路径：
- 外部调用 `store.setState`。
- store 内部 action 调用 `set`。
hydration 的核心流程是：
```text
storage.getItem(name)
-> 如果版本不一致，调用 migrate
-> merge(persistedState, currentState)
-> set(mergedState, true)
-> 触发 onFinishHydration
```
源码里还有一个 `hydrationVersion` 计数器，用于避免多次并发 `rehydrate()` 时旧请求覆盖新请求。
### `devtools`
`devtools` 连接 Redux DevTools Extension。它主要改写 `api.setState`：
```ts
api.setState = ((state, replace, nameOrAction) => {
const r = set(state, replace as any)
if (!isRecording) return r
const action =
nameOrAction === undefined
? { type: anonymousActionType || findCallerName(new Error().stack) || 'anonymous' }
: typeof nameOrAction === 'string'
? { type: nameOrAction }
: nameOrAction
connection?.send(action, get())
return r
}) as NamedSet<S>
```
所以用户可以这样给 DevTools 命名 action：
```ts
set({ count: 1 }, false, 'counter/inc')
```
它还监听 DevTools 发来的消息，支持 reset、rollback、jump to state、import state 等时间旅行能力。为了防止“从 DevTools 设置状态”再次被记录，它用 `isRecording` 临时关闭记录。
### `immer`
`immer` 改写 `store.setState`，让函数式 updater 可以直接修改 draft：
```ts
store.setState = (updater, replace, ...args) => {
const nextState = (
typeof updater === 'function' ? produce(updater as any) : updater
) as ((s: T) => T) | T | Partial<T>
return set(nextState, replace as any, ...args)
}
```
所以使用者可以写：
```ts
set((state) => {
state.deep.nested.count++
})
```
实际交给底层 `set` 的已经是 Immer 生成的不可变更新函数。
### `redux`
`redux` middleware 给 store 增加 `dispatch`：
```ts
;(api as any).dispatch = (action: A) => {
;(set as NamedSet<S>)((state: S) => reducer(state, action), false, action)
return action
}
```
它把 reducer 模式适配到 Zustand：`dispatch(action)` 本质上仍然调用 `set`，只是 next state 由 reducer 计算。
### `ssrSafe`
`ssrSafe` 是实验性 middleware。SSR 环境下它把 `setState` 替换成抛错函数：
```ts
const ssrSet = () => {
thrownew Error('Cannot set state of Zustand store in SSR')
}
api.setState = ssrSet
return config(ssrSet as never, get, api)
```
这可以阻止在服务端渲染阶段修改全局 store，避免请求之间互相污染状态。
七、类型系统如何支持 middleware
Zustand 的类型难点在 `src/vanilla.ts`：
```ts
export type StateCreator<
T,
Mis extends [StoreMutatorIdentifier, unknown][] = [],
Mos extends [StoreMutatorIdentifier, unknown][] = [],
U = T,
> = ((
setState: Get<Mutate<StoreApi<T>, Mis>, 'setState', never>,
getState: Get<Mutate<StoreApi<T>, Mis>, 'getState', never>,
store: Mutate<StoreApi<T>, Mis>,
) => U) & { $$storeMutators?: Mos }
```
可以把它理解为：
- `T`：最终 store state 类型。
- `Mis`：当前 initializer 接收到的 middleware 增强。
- `Mos`：当前 initializer 输出给外层的 middleware 增强。
- `U`：initializer 实际返回的对象类型。
`StoreMutators` 是一个空接口：
```ts
export interface StoreMutators<S, A> {}
```
middleware 通过 TypeScript module augmentation 往里面注册能力。例如 `immer`：
```ts
declare module '../vanilla' {
interface StoreMutators<S, A> {
['zustand/immer']: WithImmer<S>
}
}
```
然后 `Mutate` 根据 middleware 列表递归计算最终 store API 类型：
```ts
export type Mutate<S, Ms> =
Ms extends []
? S
: Ms extends [[infer Mi, infer Ma], ...infer Mrs]
? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
: never
```
这就是为什么用了 middleware 后，TypeScript 能知道：
- `persist` 后有 `store.persist`。
- `redux` 后有 `store.dispatch`。
- `devtools` 后 `setState` 多了 action 参数。
- `immer` 后 `setState` 支持 draft updater。
运行时靠改写 `api`，类型上靠 `StoreMutators` 注册增强，二者配合完成 middleware 机制。
八、完整调用链
以最常见用法为例：
```ts
const useCounter = create((set, get, api) => ({
count: 0,
inc: () => set((state) => ({ count: state.count + 1 })),
}))
function Counter() {
const count = useCounter((state) => state.count)
return <button>{count}</button>
}
```
调用链可以拆成两段。
创建 store：
```text
create(initializer)
-> createImpl(initializer)
-> createStore(initializer)
-> createStoreImpl(initializer)
-> 准备 setState/getState/subscribe/api
-> initializer(setState, getState, api)
-> 得到 initialState，并保存到闭包 state
-> 返回 api
-> 包装成 useBoundStore，并挂载 api 方法
```
组件订阅：
```text
useCounter(selector)
-> useStore(api, selector)
-> useSyncExternalStore(
api.subscribe,
() => selector(api.getState()),
() => selector(api.getInitialState())
)
-> React 注册订阅并读取 slice
```
状态更新：
```text
button click
-> inc()
-> set((state) => ({ count: state.count + 1 }))
-> setState 计算 nextState
-> 浅合并到 state
-> listeners.forEach(listener)
-> React 通过 useSyncExternalStore 重新读取 snapshot
-> selector 返回新的 count
-> 组件重渲染
```
九、Zustand 的设计取舍
Zustand 的源码体现了几个设计取舍：
- 小核心：核心 store 只有闭包状态、Set 订阅、set/get 四个 API。
- React 解耦：`vanilla.ts` 不依赖 React，因此可以在非 React 场景使用。
- Hook 只是绑定层：React 重渲染由 `useSyncExternalStore` 处理，不自己实现调度。
- 默认弱约束：没有强制 reducer/action，用户可以直接写函数 action。
- middleware 可组合：通过包裹 `StateCreator` 改写 API，而不是引入复杂插件系统。
- 类型增强重于运行时抽象：middleware 的很多复杂度在 TypeScript 类型层，运行时仍然很薄。
十、阅读源码建议
推荐按这个顺序读：
1. `src/vanilla.ts`：先理解 store 的完整生命周期。
2. `src/react.ts`：理解为什么 Zustand hook 能触发 React 重渲染。
3. `src/traditional.ts`、`src/vanilla/shallow.ts`、`src/react/shallow.ts`：理解 selector 和浅比较。
4. `src/middleware/subscribeWithSelector.ts`：最容易理解的 middleware。
5. `src/middleware/persist.ts`：学习持久化、hydration、异步兼容。
6. `src/middleware/devtools.ts`：学习和外部调试工具集成。
7. 回头再看 `StateCreator`、`Mutate`、`StoreMutators`：理解类型系统如何串起 middleware。
如果只记住一句话：Zustand 是一个用闭包保存状态、用 Set 管理订阅、用 `useSyncExternalStore` 接入 React、用高阶 `StateCreator` 实现 middleware 的轻量状态库。
发表评论
所有评论(0)