zustand源码解析

zustand是什么，如何使用

Zustand 是一个轻量级的 React 状态管理库，比 Redux 简单，比 Context 高效。

核心代码约 100 行，API 简单直观。
无Provider
符合 React Hooks 使用习惯
只更新订阅了变化状态的组件。
完整的类型推断和支持。

javascript 复制代码

import { create } from 'zustand'

// 1️⃣ 定义类型（TypeScript）
interface BearStore {
  bears: number
  fish: number
  addBear: () => void
  eatFish: () => void
  reset: () => void
}

// 2️⃣ 创建 Store
const useBearStore = create<BearStore>((set) => ({
  bears: 0,
  fish: 10,
  addBear: () => set((s) => ({ bears: s.bears + 1 })),
  eatFish: () => set((s) => ({ fish: s.fish - 1 })),
  reset: () => set({ bears: 0, fish: 10 })
}))

// 3️⃣ 组件 A：只订阅 bears
function BearCounter() {
  const bears = useBearStore((s) => s.bears)
  const addBear = useBearStore((s) => s.addBear)
  
  return (
    <div>
      <h2>🐻 Bears: {bears}</h2>
      <button onClick={addBear}>Add Bear</button>
    </div>
  )
}

// 4️⃣ 组件 B：只订阅 fish
function FishCounter() {
  const fish = useBearStore((s) => s.fish)
  const eatFish = useBearStore((s) => s.eatFish)
  
  return (
    <div>
      <h2>🐟 Fish: {fish}</h2>
      <button onClick={eatFish}>Eat Fish</button>
    </div>
  )
}

// 5️⃣ 主应用
function App() {
  return (
    <div>
      <BearCounter />
      <FishCounter />
    </div>
  )
}

当调用zustand的create时会发生什么

javascript 复制代码

import create from 'zustand';
const useStore = create((set) => ({
  count: 0,
  increment: () => set((s) => ({ count: s.count + 1 }))
}))

create 方法最终会调用vanvailla.ts中的createStoreImp 函数

createStoreImp主要完成的功能:

创建zustand的api(setSate,getState,getInitialState, subscribe)
在调用时会执行传入的函数，将最终结果赋值给 initialState

javascript 复制代码

// 这里的createState是下面这个函数
// (set) => ({
//   count: 0,
//   increment: () => set((s) => ({ count: s.count + 1 }))
// })
// 最终initialState的值是
// { 
//   count:0,
//   increment:() => set((s) => ({ count: s.count + 1 }))
// }

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

createStoreImp返回的api会在react.ts中的createImp中使用

javascript 复制代码

import { createStore } from './vanilla.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
}

useStore会创建一个hook, 将获取到的api 传入React.useSyncExternalStore()方法，从而实现当setSate时，订阅了这个state的组件都会触发重新渲染.

javascript 复制代码

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
}

当在一个组件调用useStore时，都会执行React.useSyncExternalStore()，useSyncExternalStore() 会执行api.subscribe. api.subscibe定义在vanilla.ts的createStoreImp函数中

javascript 复制代码

// 这里的listener是由react的useSyncExtenalStore传入的callback function, 会注册到listeners里
// 调用callback function会触发组件重新渲染
const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }
// 当 setSate里，会执行listeners里的callback function, 从而实现组件更新
 const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    // TODO: Remove type assertion once https://github.com/microsoft/TypeScript/issues/37663 is resolved
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
    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))
    }
  }

再回到react.ts中的createImpl函数，会返回一个hook给业务组件去调用

javascript 复制代码

const api = createStore(...)  // api 已创建                         │
│                                                                      │
│  // 创建 Hook 函数                                                    │
│  const useBoundStore = (selector) => useStore(api, selector)         │
│  //                                     │                            │
│  //  闭包捕获 api ◄─────────────────────┘                            │
│                                                                      │
│  // 将 api 方法挂载到 Hook 上                                         │
│  Object.assign(useBoundStore, api)                                   │
│  //                                                                  │
│  //  useBoundStore.getState = api.getState                           │
│  //  useBoundStore.setState = api.setState                           │
│  //  useBoundStore.subscribe = api.subscribe                         │
│                                                                      │
│  return useBoundStore                                                │
│

create函数主要作了以下三件事:

创建 vanilla store (包含zustand的api)
调用用户函数
包装成hook

当getState时会发生什么

javascript 复制代码

const useStore = create((set) => ({
  count: 0,
  name: 'zustand'
}))

// 获取当前状态
const state = useStore.getState()
// { count: 0, name: 'zustand' }

// 获取特定字段
const count = useStore.getState().count
// 0

根据上面的create实现解析，useStore的getState()最终会调用定义在vanilla.ts中的createStoreImpl中的getState

javascript 复制代码

const getState: StoreApi<TState>['getState'] = () => state
// 如果调用时未触发setState， 获取到的是用户在ceate时传入的state, 调用create时会执行以下代码
const initialState = (state = createState(setState, getState, api))

当setState时会发生什么

javascript 复制代码

const useStore = create((set) => ({
  count: 0
}))

// 在任何地方
useStore.setState({ count: 10 })

据上面的create实现解析，useStore的getState()最终会调用定义在vanilla.ts中的createStoreImpl中的setSate

javascript 复制代码

const setState = (partial, replace) => {
  // 1️⃣ 计算新状态
  const nextState =
    typeof partial === 'function'
      ? partial(state)    // 函数式：调用函数得到新状态
      : partial           // 对象式：直接使用

  // 2️⃣ 检查是否真的变化
  if (!Object.is(nextState, state)) {
    const previousState = state
    
    // 3️⃣ 决定合并还是替换
    state =
      (replace ?? (typeof nextState !== 'object' || nextState === null))
        ? nextState                         // 替换
        : Object.assign({}, state, nextState)  // 浅合并

    // 4️⃣ 通知所有订阅者
    listeners.forEach((listener) => listener(state, previousState))
  }
}

setSate时不会触发所有组件的更新，只会触发订阅了相关state的组件更新

javascript 复制代码

onst useStore = create((set) => ({
  count: 0,
  name: 'zustand',
  updateCount: () => set({ count: 1 }),
  updateName: () => set({ name: 'bear' })
}))

// 组件 A：只订阅 count
function CounterA() {
  console.log('CounterA 渲染')
  const count = useStore((s) => s.count)  // ← 只关心 count
  return <div>{count}</div>
}

// 组件 B：只订阅 name
function CounterB() {
  console.log('CounterB 渲染')
  const name = useStore((s) => s.name)  // ← 只关心 name
  return <div>{name}</div>
}

// 组件 C：订阅全部
function CounterC() {
  console.log('CounterC 渲染')
  const state = useStore()  // ← 关心所有
  return <div>{state.count} - {state.name}</div>
}

当组件A调用userStore时，最终会执行react.ts中的useStore

javascript 复制代码

// 这里的select就是组件A传入的(s) => s.name
export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  // getSnapshot 通过 selector 只取需要的部分, 只有需要的部分的值发送变化时，才会触发更新
  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
}

具体流程

setState 改变状态
通知所有订阅的组件（所有使用 useStore 的组件）
每个组件执行自己的 selector(newState)

比较 selector 返回的新旧值

复制代码

│

├── 相等 → 不重渲染

│

└── 不等 → 重渲染

Rect.useSyncExternalStore解析

react 18 新增的 Hook，用于安全地订阅外部数据源。

简化实现

javascript 复制代码

function useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot) {
  // 存储当前快照
  const [snapshot, setSnapshot] = useState(() => {
    // SSR 时用 getServerSnapshot
    return typeof window === 'undefined' 
      ? getServerSnapshot?.() 
      : getSnapshot()
  })

  useEffect(() => {
    // React 创建的 callback
    const callback = () => {
      const newSnapshot = getSnapshot()
      // 只有真正变化才更新
      setSnapshot(prev => {
        if (Object.is(prev, newSnapshot)) {
          return prev  // 返回旧值，不触发更新
        }
        return newSnapshot
      })
    }

    // 订阅
    const unsubscribe = subscribe(callback)

    // 清理
    return unsubscribe
  }, [subscribe, getSnapshot])

  return snapshot
}

具体流程：

Mount
├── 1.1 getSnapshot() → 初始值
├── 1.2 subscribe(callback) → 注册
└── 1.3 return 初始值
状态变化
├── 2.1 你调用 callback()
├── 2.2 React 调用 getSnapshot() → 新值
├── 2.3 Object.is(旧值, 新值)
└── 2.4 相等？
├── 2.4.1 是 → 跳过
└── 2.4.2 否 → 重渲染
Unmount
└── 3.1 调用 unsubscribe()

总结

核心实现是使用React.useSyncExternalStore hook，创建一个store，当state里的值发生改变时，触发订阅了state的值的组件重新渲染。