Zustand 核心能力与代码梳理

https://github.com/sedationh/debug-zustand

用法分析

import { create } from 'zustand'

const useStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
}))

function BearCounter() {
  const bears = useStore((state) => state.bears)
  return <h1>{bears} around here...</h1>
}

function Controls() {
  const increasePopulation = useStore((state) => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

create 传入函数，传入的函数「A」可以拿到 set 的方法用于更改状态，A 返回一个 hooks 「B」，B 可以传入 selector 来按需引用，减少渲染，B 返回所选的状态

入口

import { create } from 'zustand'

const useBearStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
}))

跳转代码 src/react.ts

export const create = (<T>(
  createState: StateCreator<T, [], []> | undefined
) => {
  // STUDY: seda 入口函数
  return createState ? createImpl(createState) : createImpl
}) as Create

进入 createImpl

import { createStore } from './vanilla.ts'

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
...

进入 createStore src/vanilla.ts

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


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) => {

可见核心能力是由 vanilla.ts 完成的，对 react 的场景在 vanilla 上包一层另外，源代码中半数以上的代码都是在写类型，如下图

考虑到这里处理的 TS 场景还比较复杂，本节不予考虑，后面单开一节写

vanilla

https://www.unpkg.com/browse/zustand@4.4.0/esm/vanilla.mjs 从官方的 mjs 版本去除一些非核心分支和判断可得下述代码

const createStoreImpl = (createState) => {
  let state
  const listeners = new Set()

  const setState = (partial, replace) => {
    const nextState = typeof partial === 'function' ? partial(state) : partial

    if (!Object.is(nextState, state)) {
      const previousState = state
      // https://docs.pmnd.rs/zustand/guides/immutable-state-and-merging#replace-flag
      // However, as this is a common pattern, set actually merges state, and we can skip the ...state part:
      state =
        replace ?? typeof nextState !== 'object'
          ? nextState
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState = () => state

  const subscribe = (listener) => {
    listeners.add(listener)
    return () => listeners.delete(listener)
  }

  const destroy = () => listeners.clear()

  const api = { setState, getState, subscribe, destroy }

  state = createState(setState, getState, api)

  return api
}
const createStore = (createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl

export { createStore }

实现了一个监听和通知用法如下，可见我们需要去手动的注册监听、取消监听、维持状态并触发更新

const vStore = createStore((set) => ({
  count: 0,
  increase: () => set((state) => ({ count: state.count + 1 })),
}))

function VanillaPage() {
  const [v, setV] = useState(() => vStore.getState())
  useEffect(() => {
    return vStore.subscribe((state) => {
      setV(state)
    })
  }, [])

  return (
    <div>
      VanillaPage
      <button
        onClick={() => {
          v.increase()
        }}>
        increase
      </button>
      {v.count}
    </div>
  )
}

react 提供一个用于订阅的 hooks — useSyncExternalStore，上面的代码可被改写为

const vStore = createStore((set) => ({
  count: 0,
  increase: () => set((state) => ({ count: state.count + 1 })),
}))

function SyncExternalStore() {
  const v = useSyncExternalStore(vStore.subscribe, vStore.getState)
  return (
    <div>
      SyncExternalStore
      <button
        onClick={() => {
          v.increase()
        }}>
        increase
      </button>
      {v.count}
    </div>
  )
}

react

结合上面的 react useSyncExternalStore 使用场景，可以理解下面的 react 实现了

import { useDebugValue } from 'react'
import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector.js'
import { createStore } from './vanilla.js'

const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports

function useStore(api, selector = api.getState, equalityFn) {
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getState,
    selector,
    equalityFn
  )
  useDebugValue(slice)
  return slice
}

const createImpl = (createState) => {
  const api = createStore(createState)

  const useBoundStore = (selector, equalityFn) =>
    useStore(api, selector, equalityFn)

  return useBoundStore
}

const create = (createState) =>
  createState ? createImpl(createState) : createImpl

export { create }

useSyncExternalStoreWithSelector 的前三个入参和 useSyncExternalStore 一样官方的相关讨论 https://github.com/reactwg/react-18/discussions/86 zustand 的的修改 https://github.com/pmndrs/zustand/pull/550/files?short_path=7ae45ad#diff-ca56e63fa839455c920562a44ebc44594f47957bbd3e9873c8a9e64104af2c41L103

之前做强制渲染的写法是

const [, forceUpdate] = useReducer((c) => c + 1, 0)

useSyncExternalStoreWithSelector 也是在 useSyncExternalStore 上进行的 memo selector 封装

selector 的能力并不是 zustand 实现的，而是交给了 react 去比对 selection 「selector 产生的 state」是否变化来决定是让组件进行 forceUpdate

middleware

如何添加一个 middleware 以 immer 为例子

import { immer } from '../mini-js/immer.js'
import { create } from '../mini-js/react.js'

export const useTodoStore = create(
  immer((set) => ({
    todos: {
      '82471c5f-4207-4b1d-abcb-b98547e01a3e': {
        id: '82471c5f-4207-4b1d-abcb-b98547e01a3e',
        title: 'Learn Zustand',
        done: false,
      },
      '354ee16c-bfdd-44d3-afa9-e93679bda367': {
        id: '354ee16c-bfdd-44d3-afa9-e93679bda367',
        title: 'Learn Jotai',
        done: false,
      },
      '771c85c5-46ea-4a11-8fed-36cc2c7be344': {
        id: '771c85c5-46ea-4a11-8fed-36cc2c7be344',
        title: 'Learn Valtio',
        done: false,
      },
      '363a4bac-083f-47f7-a0a2-aeeee153a99c': {
        id: '363a4bac-083f-47f7-a0a2-aeeee153a99c',
        title: 'Learn Signals',
        done: false,
      },
    },
    toggleTodo: (todoId) =>
      set((state) => {
        state.todos[todoId].done = !state.todos[todoId].done
      }),
  }))
)

export default () => {
  const todo = useTodoStore()
  return (
    <div>
      ImmerPage
      <h1>{JSON.stringify(todo.todos)}</h1>
      <button
        onClick={() => todo.toggleTodo('82471c5f-4207-4b1d-abcb-b98547e01a3e')}>
        toggleTodo
      </button>
    </div>
  )
}

const immer = (initializer) => (set, get, store) => {
  store.setState = (updater, replace, ...a) => {
    const nextState = typeof updater === 'function' ? produce(updater) : updater
    return set(nextState, replace, ...a)
  }
  return initializer(store.setState, get, store)
}

export { immer }

把传入的 set 方法包了一层