zustand 源码解析

实现原理

• 通过发布订阅管理数据状态，和客户端 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

SSR 数据同步

• 服务器返回的数据先存储在 localStroage 里，然后等组件 ready 后再水合 rehydrate 放进 zustand 中