Inside zustand

•5 min read
reactzustandstate managementobserver pattern
Back

zustand 參考版本為 v5.0.0-beta.0

從下方 counter 範例開始看起，主要分兩個部分：

import { create } from 'zustand'
 
type CountStore = {
  count: number
  inc: () => void
}
 
const useCountStore = create<CountStore>()((set) => ({
  count: 1,
  inc: () => set((state) => ({ count: state.count + 1 })),
}))
 
function Counter() {
  const count = useCountStore(({ count }) => count)
  const inc = useCountStore(({ inc }) => inc)
 
  return (
    <div>
      <span>{count}</span>
      <button onClick={inc}>one up</button>
    </div>
  )
}

Creating a store

第一步是透過 create，得到 useCountStore。create 對參數的判斷，可以對應到官方文件中的 TypeScript Guide，主要是為了處理型別推導的問題

The difference when using TypeScript is that instead of writing create(...), you have to write create<T>()(...) (notice the extra parentheses () too along with the type parameter) where T is the type of the state to annotate it.

💡

補充：從 discussions#1865 來看，主要與 zustand middlewares 有關


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

從上方程式碼得知，create 會調用 createImpl，這裡繼續深入 createImpl 的細節，先透過 createStore 建立一個 vanilla store

src/react.ts
import { createStore } from './vanilla.ts'
 
const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState)
  // ...
}

vanilla store 為常見的 observer pattern

The set function merges state at only one level. If you have a nested object, you need to merge them explicitly.

💡

getInitialState 主要是為了支援 server-side rendering (L26-27, L36)，相關討論在 pull#2277


src/vanilla.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) => {
    // 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))
    }
  }
 
  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

稍微理解 vanilla store 後，回到 createImpl

在第四行 useBoundStore 是接收 selector 參數的 higher order function

src/react.ts
const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState)
 
  const useBoundStore: any = (selector?: any) => useStore(api, selector)
  // ...
}

為了理解 useBoundStore 繼續看到 useStore

useStore 主要是透過 react 的 useSyncExternalStore來整合 vanilla store：

src/react.ts
export function useStore<TState, StateSlice>(
  api: StoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  const slice = useSyncExternalStore(
    api.subscribe,
    () => selector(api.getState()),
    () => selector(api.getInitialState()),
  )
  useDebugValue(slice)
  return slice
}

再回到 createImpl 的程式碼，整理一下目前已知：

src/react.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
}
💡

補充：Object.assign(useBoundStore, api)，可以對應到官方文件中提到：

Reading/writig state and reacting to changes outside of components

const useDogStore = create(() => ({ paw: true, snout: true, fur: true }))
 
// Getting non-reactive fresh state
const paw = useDogStore.getState().paw
// Listening to all changes, fires synchronously on every change
const unsub = useDogStore.subscribe(console.log)
useDogStore.setState({ paw: false })
// Unsubscribe listeners
unsub()

Binding Components

在上述結果得知 useBoundStore 是基於 useSyncExternalStore 實現，當調用 useBoundStore(selector)：

While the store has not changed, repeated calls to getSnapshot must return the same value. If the store changes and the returned value is different (as compared by Object.is), React re-renders the component.

💡

補充：在目前 react 的版本（v18.3.1），任何同步外部儲存的狀態，都需要透過 useSyncExternalStore 來避免 tearing。除了 zustand 外，@tanstack/react-query 也可以看到 useSyncExternalStore 的身影

Summary

以下簡單總結：

Reference