likes
comments
collection
share

手写状态管理库？一篇文章带你深入了解Zustand！

作者站长头像
站长
· 阅读数 58

Zustand：一个轻量，简洁的状态管理库。它可以与技术栈无关，可以在React中直接使用，也可以在Vue中直接调用内置的api。它的源码也十分简单，非常会借力，大致思路是创建一个闭包函数用来存储全局数据，通过发布-订阅模式来触发组件的重新渲染，这也就是为什么它可以在无React依赖的情况下使用，我们可以自己实现监听器来出触发reRender。下面一起来看下这个有趣的状态管理库吧！

创建一个Store

我们通过调用create方式创建了一个hooks函数，store存储的数据可以是基本数据类型（string number boolean null undefined）, 也可以是object 和 function。我们可以function通过set函数来合并我们的store。

import { create } from 'zustand'

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

在组件中使用

我们可以在组件中直接调用创建的hooks，入参传递函数来选择store中我们需要的值，hooks 可以直接将单个数据返回，也可以将多个数据做成对象或数组的形式返回。也可以选择什么都不传，那么就是取的整个store里的值，任何数据更新都会造成重新渲染！

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

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

我们来控制数据更新是否需要重新render？

Zustand的默认的对比策略是全等判断（===），但是如果数据以对象或数组的形式返回，即使数据内容相同，由于前后不是同一个对象，依然会触发组件重绘。zustand 提供了 shallow 函数，用来比较对象或数组，如果浅比较值不相等的话才会更新，作为第二个参数可以解决多次触发重绘的问题。也可以传入一个定义函数来处理是否需要更新，如果返回true则表示数据相等不需要更新。

import { shallow } from 'zustand/shallow'

// Object pick, re-renders the component when either state.nuts or state.honey change
const { nuts, honey } = useBearStore(
  (state) => ({ nuts: state.nuts, honey: state.honey }),
  shallow
)

// Array pick, re-renders the component when either state.nuts or state.honey change
const [nuts, honey] = useBearStore(
  (state) => [state.nuts, state.honey],
  shallow
)

// Mapped picks, re-renders the component when state.treats changes in order, count or keys
const treats = useBearStore((state) => Object.keys(state.treats), shallow)

// 自定义比较函数
const treats = useBearStore(
  (state) => state.treats,
  (oldTreats, newTreats) => compare(oldTreats, newTreats)
)

Set 的第二个参数 replace

set函数更新的时候每次都会新生成一个store，新的值会替换旧的值；默认为false，会浅拷贝合并到旧的store中，我们可以手动传入true则替换整个store；

set((store) => ({ count: store.count + 1 })) 
// == 相当于
set((store) => ({ ...store, count: store.count + 1 }))


import omit from 'lodash-es/omit'

const useFishStore = create((set) => ({
  salmon: 1,
  tuna: 2,
  deleteEverything: () => set({}, true), //  会情况store中的全部数据
  deleteTuna: () => set((state) => omit(state, ['tuna']), true), // 替换最新数据
}))

组件外部使用Store的数据

生成的hooks在原型上提供一些store内部api，如 getState setState 或 subscribe；这些Api极大为Zustand提供了极大的扩展性，比如我们可以在非React组件中去更新或获取Store里的值...

const useDogStore = create(() => ({ paw: true, snout: true, fur: true }))

// Getting non-reactive fresh state
const paw = useDogStore.getState().paw
// Listening to all changes, fires synchronously on every change
const unsub1 = useDogStore.subscribe(console.log)
// Updating state, will trigger listeners
useDogStore.setState({ paw: false })
// Unsubscribe listeners
unsub1()

// You can of course use the hook as you always would
const Component = () => {
  const paw = useDogStore((state) => state.paw)
  ...

中间件

zustand也提供了一些类似于Redux的中间件，例如log，immer，devtools等，所有的中间件都是可以一起使用的，中间件的存在也一方面提供了Zustand的丰富性和可扩展。

// persist中间件，localStorage本地缓存数据
import { create } from 'zustand'
import { persist, createJSONStorage } from 'zustand/middleware'

const useFishStore = create(
  persist(
    (set, get) => ({
      fishes: 0,
      addAFish: () => set({ fishes: get().fishes + 1 }),
    }),
    {
      name: 'food-storage', // unique name
      storage: createJSONStorage(() => sessionStorage), // (optional) by default, 'localStorage' is used
    }
  )
)

一起来看源码吧！

看一个项目的源码，我们首先要看它的目录文件，先看下项目的整体结构，找到它的入口文件从项目的入口入手，我们调用zustand的时候，首先是用的create函数，那么我们就从代码库中先找到create函数。

  1. index.ts作为项目入口，引入了vanilla和react2个文件
  2. react.ts 调用 vanilla 中的createStore并获取api，并注册监听处理函数等。
  3. vanilla.ts 核心代码，维护store数据和监听器集合，并提供api方法暴露给消费者
  4. shallow.ts 提供浅比较函数给hooks的第二个参数使用
  5. middleware 中间件文件夹，提供了一些内置的中间件

vanilla.ts

先贴下当前版本源码

type SetStateInternal<T> = {
  _(
    partial: T | Partial<T> | { _(state: T): T | Partial<T> }['_'],
    replace?: boolean | undefined
  ): void
}['_']

export interface StoreApi<T> {
  setState: SetStateInternal<T>
  getState: () => T
  subscribe: (listener: (state: T, prevState: T) => void) => () => void
  /**
   * @deprecated Use `unsubscribe` returned by `subscribe`
   */
  destroy: () => void
}

type Get<T, K, F> = K extends keyof T ? T[K] : F

export type Mutate<S, Ms> = number extends Ms['length' & keyof Ms]
  ? S
  : Ms extends []
  ? S
  : Ms extends [[infer Mi, infer Ma], ...infer Mrs]
  ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
  : never

export type StateCreator<
  T,
  Mis extends [StoreMutatorIdentifier, unknown][] = [],
  Mos extends [StoreMutatorIdentifier, unknown][] = [],
  U = T
> = ((
  setState: Get<Mutate<StoreApi<T>, Mis>, 'setState', never>,
  getState: Get<Mutate<StoreApi<T>, Mis>, 'getState', never>,
  store: Mutate<StoreApi<T>, Mis>
) => U) & { $$storeMutators?: Mos }

// eslint-disable-next-line @typescript-eslint/no-unused-vars, @typescript-eslint/no-empty-interface
export interface StoreMutators<S, A> {}
export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>

type CreateStore = {
  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>
  ): Mutate<StoreApi<T>, Mos>

  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>
  ) => Mutate<StoreApi<T>, Mos>
}

type CreateStoreImpl = <
  T,
  Mos extends [StoreMutatorIdentifier, unknown][] = []
>(
  initializer: StateCreator<T, [], Mos>
) => Mutate<StoreApi<T>, Mos>

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) => {
    // TODO: Remove type assertion once https://github.com/microsoft/TypeScript/issues/37663 is resolved
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        replace ?? typeof nextState !== 'object'
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState: StoreApi<TState>['getState'] = () => state

  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }

  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 }
  state = createState(setState, getState, api)
  return api as any
}

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

/**
 * @deprecated Use `import { createStore } from 'zustand/vanilla'`
 */
export default ((createState) => {
  if (import.meta.env?.MODE !== 'production') {
    console.warn(
      "[DEPRECATED] Default export is deprecated. Instead use import { createStore } from 'zustand/vanilla'."
    )
  }
  return createStore(createState)
}) as CreateStore

// ---------------------------------------------------------

/**
 * @deprecated Use `unknown` instead of `State`
 */
export type State = unknown

/**
 * @deprecated Use `Partial<T> | ((s: T) => Partial<T>)` instead of `PartialState<T>`
 */
export type PartialState<T extends State> =
  | Partial<T>
  | ((state: T) => Partial<T>)

/**
 * @deprecated Use `(s: T) => U` instead of `StateSelector<T, U>`
 */
export type StateSelector<T extends State, U> = (state: T) => U

/**
 * @deprecated Use `(a: T, b: T) => boolean` instead of `EqualityChecker<T>`
 */
export type EqualityChecker<T> = (state: T, newState: T) => boolean

/**
 * @deprecated Use `(state: T, previousState: T) => void` instead of `StateListener<T>`
 */
export type StateListener<T> = (state: T, previousState: T) => void

/**
 * @deprecated Use `(slice: T, previousSlice: T) => void` instead of `StateSliceListener<T>`.
 */
export type StateSliceListener<T> = (slice: T, previousSlice: T) => void

/**
 * @deprecated Use `(listener: (state: T) => void) => void` instead of `Subscribe<T>`.
 */
export type Subscribe<T extends State> = {
  (listener: (state: T, previousState: T) => void): () => void
}

/**
 * @deprecated You might be looking for `StateCreator`, if not then
 * use `StoreApi<T>['setState']` instead of `SetState<T>`.
 */
export type SetState<T extends State> = {
  _(
    partial: T | Partial<T> | { _(state: T): T | Partial<T> }['_'],
    replace?: boolean | undefined
  ): void
}['_']

/**
 * @deprecated You might be looking for `StateCreator`, if not then
 * use `StoreApi<T>['getState']` instead of `GetState<T>`.
 */
export type GetState<T extends State> = () => T

/**
 * @deprecated Use `StoreApi<T>['destroy']` instead of `Destroy`.
 */
export type Destroy = () => void

1~59 在定义下面的一些类型，看ts类型定义有助了解代码哦，这里面用了很多泛型和函数重载，大家可以学习一下。

60~103 createStoreImpl函数为创建store的核心代码

63~64 定义全局的state和监听器list

66~81 setState用来更新数据，接收2个入参；第一个入参可能为函数或者直接为值，我们取到最新值和旧的值比较一下，看下有没有更新；如果有更新的话，根据第二个参数replace来判断是进行合并还是直接代替。把最新的值赋值给state，并执行监听器list中的全部监听函数。

83 getState直接返回state

85~89 subscribe 接收监听函数，并添加到监听器list中，并在销毁时从监听器list中清除；等待setState更新state后调用；

91~ 98 清除监听器list

100 提供api方法，通过102行return供外部使用

101 通过createStoreImpl函数的函数参数createState把set，get，api三个方法暴露出去，然后初始化state，set和get 我们可以在注册hooks的时候使用，api大部分用来提供给中间件使用。

这样整体就产生了一个闭包，createStoreImpl函数虽然调用完了，但是state并不会销毁。再后面也就是一些类型定义了。

react.ts

贴下代码,那我先从调用顺序开始看

import { useDebugValue } from 'react'
// import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector'
// This doesn't work in ESM, because use-sync-external-store only exposes CJS.
// See: https://github.com/pmndrs/valtio/issues/452
// The following is a workaround until ESM is supported.
// eslint-disable-next-line import/extensions
import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector'
import { createStore } from './vanilla.ts'
import type {
  Mutate,
  StateCreator,
  StoreApi,
  StoreMutatorIdentifier,
} from './vanilla.ts'

const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports

type ExtractState<S> = S extends { getState: () => infer T } ? T : never

type ReadonlyStoreApi<T> = Pick<StoreApi<T>, 'getState' | 'subscribe'>

type WithReact<S extends ReadonlyStoreApi<unknown>> = S & {
  getServerState?: () => ExtractState<S>
}

export function useStore<S extends WithReact<StoreApi<unknown>>>(
  api: S
): ExtractState<S>

export function useStore<S extends WithReact<StoreApi<unknown>>, U>(
  api: S,
  selector: (state: ExtractState<S>) => U,
  equalityFn?: (a: U, b: U) => boolean
): U

export function useStore<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
}

export type UseBoundStore<S extends WithReact<ReadonlyStoreApi<unknown>>> = {
  (): ExtractState<S>
  <U>(
    selector: (state: ExtractState<S>) => U,
    equals?: (a: U, b: U) => boolean
  ): U
} & S

type Create = {
  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>
  ): UseBoundStore<Mutate<StoreApi<T>, Mos>>
  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>
  ) => UseBoundStore<Mutate<StoreApi<T>, Mos>>
  /**
   * @deprecated Use `useStore` hook to bind store
   */
  <S extends StoreApi<unknown>>(store: S): UseBoundStore<S>
}

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

  const useBoundStore: any = (selector?: any, equalityFn?: any) =>
    useStore(api, selector, equalityFn)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

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

/**
 * @deprecated Use `import { create } from 'zustand'`
 */
export default ((createState: any) => {
  if (import.meta.env?.MODE !== 'production') {
    console.warn(
      "[DEPRECATED] Default export is deprecated. Instead use `import { create } from 'zustand'`."
    )
  }
  return create(createState)
}) as Create

93~94 就是我们调用create方法，create又调用了createImpl；

73~91 createImpl函数，限制了create函数的入参必须也是个函数；

82~83 使用vanilla.ts中的createStore初始化store，并返回store中的相关api

85~86 通过createStore返回我们需要内容 在88行给useBoundStore添加api方法

36~50 useStore就是在react内更新组件核心方法，也是这个库很会借力的地方直接调用useSyncExternalStoreWithSelector,把更新兼容组件更新的操作交给react来做，如果要做一个vue版本的zustand那么替换useStore为vue相关的代码就可以了。我们也可以用useReducer来代理 useSyncExternalStoreWithSelector来强制更新就可以了。

  • useSyncExternalStore是一个自定义 Hook，它提供了一种简单的方法来订阅外部状态管理器，并将其状态同步到 React 组件中
  • useSyncExternalStore函数的第一个参数是一个订阅函数，它接收一个回调函数作为参数，当状态发生变化时，该回调函数将被调用。该回调函数接收两个参数：当前的状态值和上一个状态值
  • useSyncExternalStore函数的第二个参数是一个获取状态值的函数。当组件需要获取当前状态值时，它将调用该函数并返回当前状态值
  • useSyncExternalStore 函数返回一个状态值，该值表示当前的状态。当外部状态管理器更新状态时，组件将自动更新状态

中间件

以immer为例：倒也是比较借力，直接用的immer的produce方法，可以解决一些对象层级过深的问题。

const immerImpl: ImmerImpl = (initializer) => (set, get, store) => {
  type T = ReturnType<typeof initializer>

  store.setState = (updater, replace, ...a) => {
    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

在中间件内先重写setState，再调用zustand内置的setState；

Vue-Zustand

利用vue的ref的响应式，创建一个ref对象就可以了，监听器数据更新的时候重新给ref对象赋值。

import type { UnwrapRef } from 'vue'
import { getCurrentInstance, onScopeDispose, readonly, ref, toRefs } from 'vue'
import { toReactive } from '@vueuse/core'

import type {
  Mutate,
  StateCreator,
  StoreApi,
  StoreMutatorIdentifier,
} from 'zustand/vanilla'
import { createStore as createZustandStore } from 'zustand/vanilla'
import type { IsPrimitive } from './util'
import { isPrimitive } from './util'

type ExtractState<S> = S extends { getState: () => infer T } ? T : never

export function useStore<S extends StoreApi<unknown>>(api: S): ExtractState<S>

export function useStore<S extends StoreApi<unknown>, U>(
  api: S,
  selector: (state: ExtractState<S>) => U,
  equalityFn?: (a: U, b: U) => boolean
): U

export function useStore<TState extends object, StateSlice>(
  api: StoreApi<TState>,
  selector: (state: TState) => StateSlice = api.getState as any,
  equalityFn?: (a: StateSlice, b: StateSlice) => boolean,
) {
  const initialValue = selector(api.getState())
  const state = ref(initialValue)

  const listener = (nextState: TState, previousState: TState) => {
    const prevStateSlice = selector(previousState)
    const nextStateSlice = selector(nextState)

    if (equalityFn !== undefined) {
      if (!equalityFn(prevStateSlice, nextStateSlice))
        state.value = nextStateSlice as UnwrapRef<StateSlice>
    }
    else {
      state.value = nextStateSlice as UnwrapRef<StateSlice>
    }
  }

  const unsubscribe = api.subscribe(listener)

  if (getCurrentInstance()) {
    onScopeDispose(() => {
      unsubscribe()
    })
  }

  return isPrimitive(state.value) ? readonly(state) : toRefs(toReactive(state))
}

export type UseBoundStore<S extends StoreApi<unknown>> = {
  (): IsPrimitive<ExtractState<S>>
  <U>(
    selector: (state: ExtractState<S>) => U,
    equals?: (a: U, b: U) => boolean
  ): IsPrimitive<U>
} & S

interface Create {
  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>
  ): UseBoundStore<Mutate<StoreApi<T>, Mos>>
  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>
  ) => UseBoundStore<Mutate<StoreApi<T>, Mos>>
  <S extends StoreApi<unknown>>(store: S): UseBoundStore<S>
}

const createImpl = <T extends object>(createState: StateCreator<T, [], []>) => {
  const api
    = typeof createState === 'function' ? createZustandStore(createState) : createState

  const useBoundStore: any = (selector?: any, equalityFn?: any) =>
    useStore(api, selector, equalityFn)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

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

export default create

demo

简单仿官网搞个一个纯js版的zustand，可以通过/react或者/vue来使用不同框架对应的zustand；也有提供middleware和shallow方法哦！

vue3版demo：codesandbox.io/p/sandbox/v…

react版demo：codesandbox.io/p/sandbox/r…

npm包：www.npmjs.com/package/toy…

git地址： github.com/waltiu/mini…

转载自:https://juejin.cn/post/7254097346761998393
评论
请登录