zustand 源碼解析

文章目錄

實現原理
create
createStore 創建實例
CreateStoreImpl 實現發布訂閱
createImpl 包裝返回給用戶調用的 hook
useSyncExternalStoreWithSelector 訂閱更新
zustand 性能優化
- 自定義數據更新
- - createWithEqualityFn
  - createWithEqualityFnImpl 返回 hook
  - useSyncExternalStoreWithSelector 自定義比較函數
- 使用 useShallow 淺比較
- - useShallow
  - shallow 淺比較實現
- 使用 immer
- - 在組件外部數據訂閱
SSR 數據同步

實現原理

通過發布訂閱管理數據狀態，和客戶端 useSyncExternalStoreWithSelector 解耦，支持 SSR
通過 useSyncExternalStoreWithSelector 結合上面步驟實現數據更新
通過 useSyncExternalStoreWithSelector 的自定義切片數據、自定義是否更新狀態函數、淺比較數據實現性能優化
數據 set 時通過Object.is 比較是否變化
- 相比 === 的差異在 NAN、±0 的比較上
數據 set 時的合并通過 Object.assign 實現

create

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

createStore 創建實例

export const createStore = ((createState) =>createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

CreateStoreImpl 實現發布訂閱

通過閉包保存 state
通過發布訂閱管理 state
數據的合并通過 Object.assign實現
這樣實現的好處是可以在服務端使用，而不依賴于 useSyncExternalStoreWithSelector

const createStoreImpl: CreateStoreImpl = (createState) => {type TState = ReturnType<typeof createState>type Listener = (state: TState, prevState: TState) => voidlet state: TState// Set 去重const listeners: Set<Listener> = new Set()const setState: StoreApi<TState>['setState'] = (partial, replace) => {// 判斷 set 時傳入的是對象還是函數const nextState =typeof partial === 'function'? (partial as (state: TState) => TState)(state): partial// 判斷數據是否變化if (!Object.is(nextState, state)) {const previousState = state// 是否直接 replace 替換 state，還是 merge 合并 statestate =replace ?? (typeof nextState !== 'object' || nextState === null)? (nextState as TState): Object.assign({}, state, nextState)// 數據變換后發布通知listeners.forEach((listener) => listener(state, previousState))}}// 直接獲取保存的閉包 stateconst getState: StoreApi<TState>['getState'] = () => state// 添加訂閱const subscribe: StoreApi<TState>['subscribe'] = (listener) => {listeners.add(listener)// Unsubscribereturn () => listeners.delete(listener)}// 取消所有訂閱直接調用 Set 的 clearconst destroy: StoreApi<TState>['destroy'] = () => {if (import.meta.env?.MODE !== 'production') {console.warn('[DEPRECATED] The `destroy` method will be unsupported in a future version. Instead use unsubscribe function returned by subscribe. Everything will be garbage-collected if store is garbage-collected.',)}listeners.clear()}const api = { setState, getState, subscribe, destroy }// createState 是用戶傳入的函數,傳入的函數會拿到 setState、getState 和 api 方法state = createState(setState, getState, api)return api as any
}

createImpl 包裝返回給用戶調用的 hook

實際是先調用上面的 createStore 先為狀態創建對應的發布訂閱實例
再通過 useStore 將發布訂閱和 useSyncExternalStoreWithSelector 聯系起來進而可以發布訂閱后通知 React 進行更新

const createImpl = <T>(createState: StateCreator<T, [], []>) => {if (import.meta.env?.MODE !== 'production' &&typeof createState !== 'function') {console.warn("[DEPRECATED] Passing a vanilla store will be unsupported in a future version. Instead use `import { useStore } from 'zustand'`.",)}// 管理狀態的實例（發布訂閱）const api =typeof createState === 'function' ? createStore(createState) : createState// 交給用戶調用的 hook，通過 useStore 封裝 useSyncExternalStoreWithSelectorconst useBoundStore: any = (selector?: any, equalityFn?: any) =>useStore(api, selector, equalityFn)// 添加額外 APIObject.assign(useBoundStore, api)return useBoundStore
}

useSyncExternalStoreWithSelector 訂閱更新

useSyncExternalStoreWithSelector 比 useSyncExternalStore 性能更好，可以定義想訂閱的狀態

// 通過 useSyncExternalStoreWithSelector hook 訂閱外部存儲，并返回你關心的數據部分
export function useStore<TState, StateSlice>(api: WithReact<StoreApi<TState>>,selector: (state: TState) => StateSlice = api.getState as any,equalityFn?: (a: StateSlice, b: StateSlice) => boolean,
) {if (import.meta.env?.MODE !== 'production' &&equalityFn &&!didWarnAboutEqualityFn) {console.warn("[DEPRECATED] Use `createWithEqualityFn` instead of `create` or use `useStoreWithEqualityFn` instead of `useStore`. They can be imported from 'zustand/traditional'. https://github.com/pmndrs/zustand/discussions/1937",)didWarnAboutEqualityFn = true}// 返回狀態，此時發布訂閱和 React 更新關聯上了const slice = useSyncExternalStoreWithSelector(api.subscribe,api.getState,api.getServerState || api.getState,selector, // 選擇外部存儲中你關心的數據部分equalityFn,// 自定義是否重新渲染)useDebugValue(slice)return slice
}

zustand 性能優化

自定義數據更新

zustand 支持傳入一個自定義比較函數確定數據是否變化
- 實質是 useSyncExternalStoreWithSelector 支持傳入比較函數

createWithEqualityFn

不適用默認的 create ，使用 createWithEqualityFn 并傳入比較函數

export const createWithEqualityFn = (<T>(createState: StateCreator<T, [], []> | undefined,defaultEqualityFn?: <U>(a: U, b: U) => boolean,
) =>createState? createWithEqualityFnImpl(createState, defaultEqualityFn): createWithEqualityFnImpl) as CreateWithEqualityFn

createWithEqualityFnImpl 返回 hook

取消 Object.is 比較數據變化

const createWithEqualityFnImpl = <T>(createState: StateCreator<T, [], []>,defaultEqualityFn?: <U>(a: U, b: U) => boolean,
) => {// 創建發布訂閱實例const api = createStore(createState)// 傳入自定義比較函數，取消默認的 Object.is 比較const useBoundStoreWithEqualityFn: any = (selector?: any,equalityFn = defaultEqualityFn,) => useStoreWithEqualityFn(api, selector, equalityFn)Object.assign(useBoundStoreWithEqualityFn, api)// 返回用戶調用的 hook return useBoundStoreWithEqualityFn
}

useSyncExternalStoreWithSelector 自定義比較函數

export function useStoreWithEqualityFn<S extends WithReact<StoreApi<unknown>>,U,
>(api: S,selector: (state: ExtractState<S>) => U,equalityFn?: (a: U, b: U) => boolean,
): Uexport function useStoreWithEqualityFn<TState, StateSlice>(api: WithReact<StoreApi<TState>>,selector: (state: TState) => StateSlice = api.getState as any,equalityFn?: (a: StateSlice, b: StateSlice) => boolean,
) {const slice = useSyncExternalStoreWithSelector(api.subscribe,api.getState,api.getServerState || api.getState,selector,equalityFn,)useDebugValue(slice)return slice
}

使用 useShallow 淺比較

可以通過 useShallow 淺比較提升性能

const { nuts, honey } = useBearStore(useShallow((state) => ({ nuts: state.nuts, honey: state.honey })),
)

淺比較的更新策略，可以看到相同屬性的對象就算是新對象，比較時都會默認無變化提升性能

  shallow(1, 1); // trueshallow('hello', 'hello'); // trueshallow({ a: 1 }, { a: 1 }); // trueshallow([1, 2, 3], [1, 2, 3]); // false

useShallow

保存 selector 結果，對應的結果會給到 useSyncExternalStoreWithSelector 的第四個參數，如果提供了第五個參數比較函數，則使用比較函數判斷兩次結果，默認第五個比較函數為 Object.is ，會跳過更新

  const slice = useSyncExternalStoreWithSelector(api.subscribe,api.getState,api.getServerState || api.getState,  selector, // 選擇外部存儲中你關心的數據部分，對應 useShallow 返回的函數equalityFn,// 自定義是否重新渲染，默認為Object.is)

const { useRef } = ReactExportsexport function useShallow<S, U>(selector: (state: S) => U): (state: S) => U {const prev = useRef<U>()// 通過閉包保存 selector 返回的結果，對應的結果會給到 useSyncExternalStoreWithSelector 的第四個參數 selector，如果返回的引用不變，不會出觸發更新return (state) => {const next = selector(state)return shallow(prev.current, next)? (prev.current as U): (prev.current = next)}
}

shallow 淺比較實現

最主要的是判斷當對象 key-value 一樣時，不管是否是新建對象都會是認為一樣的

export function shallow<T>(objA: T, objB: T) {if (Object.is(objA, objB)) {return true}if (typeof objA !== 'object' ||objA === null ||typeof objB !== 'object' ||objB === null) {return false}if (objA instanceof Map && objB instanceof Map) {if (objA.size !== objB.size) return falsefor (const [key, value] of objA) {if (!Object.is(value, objB.get(key))) {return false}}return true}if (objA instanceof Set && objB instanceof Set) {if (objA.size !== objB.size) return falsefor (const value of objA) {if (!objB.has(value)) {return false}}return true}const keysA = Object.keys(objA)if (keysA.length !== Object.keys(objB).length) {return false}for (let i = 0; i < keysA.length; i++) {if (!Object.prototype.hasOwnProperty.call(objB, keysA[i] as string) ||!Object.is(objA[keysA[i] as keyof T], objB[keysA[i] as keyof T])) {return false}}return true
}/*** shallow(1, 1); // trueshallow('hello', 'hello'); // trueshallow({ a: 1 }, { a: 1 }); // trueshallow([1, 2, 3], [1, 2, 3]); // false*/

使用 immer

immer 的中間件實現主要是攔截了發布訂閱的 set 方法，通過immer去修改數據后，再調用原始 set 方法

const useBeeStore = create(immer((set) => ({bees: 0,addBees: (by) =>set((state) => {state.bees += by}),})),
)

const immerImpl: ImmerImpl = (initializer) => (set, get, store) => {type T = ReturnType<typeof initializer>// 攔截了 set 方法 store.setState = (updater, replace, ...a) => {// 通過 immer 修改數據const nextState = (typeof updater === 'function' ? produce(updater as any) : updater) as ((s: T) => T) | T | Partial<T>// 修改后再調用原始方法return set(nextState as any, replace, ...a)}return initializer(store.setState, get, store)
}export const immer = immerImpl as unknown as Immer

在組件外部數據訂閱

同樣是攔截了 subscribe 方法，添加訂閱邏輯

const subscribeWithSelectorImpl: SubscribeWithSelectorImpl =(fn) => (set, get, api) => {type S = ReturnType<typeof fn>type Listener = (state: S, previousState: S) => voidconst origSubscribe = api.subscribe as (listener: Listener) => () => void// 攔截 subscribe api.subscribe = ((selector: any, optListener: any, options: any) => {let listener: Listener = selector // if no selectorif (optListener) {const equalityFn = options?.equalityFn || Object.is// TODO：這樣后續更新后，無論如何都會更新，因為只比較第一次保存的狀態let currentSlice = selector(api.getState())// 添加本次 listenerlistener = (state) => {const nextSlice = selector(state)// 只有當數據變化時才觸發if (!equalityFn(currentSlice, nextSlice)) {const previousSlice = currentSliceoptListener((currentSlice = nextSlice), previousSlice)}}if (options?.fireImmediately) {optListener(currentSlice, currentSlice)}}return origSubscribe(listener)}) as anyconst initialState = fn(set, get, api)return initialState}
export const subscribeWithSelector =subscribeWithSelectorImpl as unknown as SubscribeWithSelector