从零实现一个 zustand

使用

Zustand 使用起来非常简单，使用 create 创建一个 store 后，既可在全局使用 store 中的值

tsx 复制代码

import { create } from 'zustand'

// 创建 store
const useStore = create((set) => ({
  count: 1,
  inc: () => set((state) => ({ count: state.count + 1 })),
}))


// 组件中使用
function Counter() {
  const count = useStore((state) => state.count)
  const inc = useStore((state) => state.inc)

  return (
    <div className="counter">
      <span>{count}</span>
      <button onClick={inc}>one up</button>
    </div>
  )
}

注意：当我们在组件中利用下面这种方式使用 store 中的某个值

tsx 复制代码

// 只要 store 中的任意一个值变化，都会导致使用 inc 和 count 这个值的组件重新渲染
const { inc, count } = useStore();

只要 store 中的某个值变化，组件都会进行重新渲染

所以，我们可以利用 selector 或者 shallow 方式来进行取值

selector

tsx 复制代码

 const count = useStore((state) => state.count)

shallow

tsx 复制代码

import { useShallow } from 'zustand/shallow';

const { inc, count } = useStore(
  useShallow((state) => ({
    inc: state.inc,
    count: state.count,
  })),
);

Zustand 还提供了一些中间件 immer 等，使用也非常简单。具体使用可以参考官网

源码

zustand 实现非常简单，主要是利用了发布订阅的设计模式，在 React 中，通过 useSyncExternalStore 这个 hooks 来实现渲染组件

Zustand 利用 create 创建了一个 store 进行返回, 并且返回的这个 store 是一个函数，以供我们通过 useStore 的形式调用

create 创建 store 的时候，create 的参数也是一个函数，他会将我们的状态值进行传入后创建。利用 useSyncExternalStore 来更新

ts 复制代码

export const create = (createState) => {
  
  const api = createStore(createState)
  
  return () => useSyncExternalStore(
    api.subscribe,
    api.getState,
  )
}

下面来实现 createStore 函数直接一个发布订阅，createState 是一个函数，这个函数接收 setState, getState, state 参数然后返回 state

ts 复制代码

export const createStore = (createState: any) => {
  type TState = ReturnType<typeof createState>
  type Listener = (state: TState, prevState: TState) => void

  let state: TState
  let listeners: Set<Listener> = new Set()

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

    if (!Object.is(state, nextState)) {
      const previousState = state

      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? nextState
          : Object.assign({}, state, nextState)

      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState = () => state

  const subscribe = (listener) => {
    listeners.add(listener)

    return () => listener.delete(listener)
  }

  const getInitialState = () => initialState

  const initialState = createState(setState, getState, state)

  state = initialState

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

  return api
}

然后我们实现一下 selector 的使用方法，修改一下 create 函数

diff 复制代码

const useStore = (api, selector) => {

  const slice = useSyncExternalStore(
    api.subscribe,
    // api.getState,
    selector ? () => selector(api.getState()) : api.getState,
  )

  return slice
}

export const create = (createState: any) => {
  const api = createStore(createState) 
  
  const useBoundStore: any = (selector?: any) => useStore(api, selector)
  Object.assign(useBoundStore, api)

  return useBoundStore
}

大功告成！！