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) => void
  
  let 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 合并 state
      state =
        replace ?? (typeof nextState !== 'object' || nextState === null)
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      // 数据变换后发布通知
      listeners.forEach((listener) => listener(state, previousState))
    }
  }
  
  // 直接获取保存的闭包 state
  const getState: StoreApi<TState>['getState'] = () => state
  // 添加订阅
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }
  
  // 取消所有订阅直接调用 Set 的 clear
  const 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 封装 useSyncExternalStoreWithSelector
  const useBoundStore: any = (selector?: any, equalityFn?: any) =>
    useStore(api, selector, equalityFn)

  // 添加额外 API
  Object.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,
): U

export 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); // true
  shallow('hello', 'hello'); // true
  shallow({ a: 1 }, { a: 1 }); // true
  shallow([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 } = ReactExports

export 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 false

    for (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 false

    for (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); // true
  shallow('hello', 'hello'); // true
  shallow({ a: 1 }, { a: 1 }); // true
  shallow([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) => void
    const origSubscribe = api.subscribe as (listener: Listener) => () => void
    // 拦截 subscribe 
    api.subscribe = ((selector: any, optListener: any, options: any) => {
      let listener: Listener = selector // if no selector
      if (optListener) {
        const equalityFn = options?.equalityFn || Object.is
        // TODO：这样后续更新后，无论如何都会更新，因为只比较第一次保存的状态
        let currentSlice = selector(api.getState())
        // 添加本次 listener
        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
    const initialState = fn(set, get, api)
    return initialState
  }
export const subscribeWithSelector =
  subscribeWithSelectorImpl as unknown as SubscribeWithSelector