【源码共读】| 简易实现zustand

1,567 阅读4分钟

Zustand是一个函数式的状态管理库,因为它简单易用且上手容易,所以社区对其非常热衷。

本文主要解决以下几个问题:

  • 与Redux的差异?
  • 源码分析
  • 源码的简单实现

image.png


首先,我们来看下Zustand的用法

简单使用

那么,我们先来看看他的使用:

docs.pmnd.rs/zustand/get…

// 定义仓库
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>
}

与Redux的对比

特性ReduxZustand
状态管理全局状态管理库,适用于大型应用轻量级状态管理库,适用于小型到中型应用
模式严格的模式,需要定义actions和reducers灵活的模式,可以直接在store中定义函数
中间件支持中间件,可以在action被发送到store之前或之后进行操作可以通过hooks的形式实现自定义中间件,内置的也有immer和数据持久化
开发工具有强大的开发者工具,可以查看状态的历史记录,进行时间旅行调试等没有像Redux那样的开发者工具,但可以通过React Devtools进行调试
社区支持有大量的社区支持和丰富的插件库社区支持较少,但由于其简洁的API,学习曲线较低
优点严格的模式有助于保持代码的一致性和可维护性;强大的开发工具可以提高开发效率API简洁,学习曲线低;更灵活的模式可以更自由地组织代码
缺点严格的模式可能导致代码冗余;对于小型应用来说可能过于复杂

源码分析

那么,他是如何通过简洁的代码实现状态管理的呢?让我们来看一下他的源码:github.com/pmndrs/zust…
这里摘取了部分来说明:

  • 使用createStore包装成api
  • 借助useSyncExternalStoreWithSelector来实现数据更新
  • 使用了发布订阅的数据更新模式

image.png
image.png
image.png
Zustand是一个基于React Hooks的轻量级状态管理库,其原理是通过创建全局的可观察状态,然后在组件中使用hooks订阅这些状态,当状态改变时,相关的组件会自动重新渲染。


简易实现

首先,我们来实现vanilla的部分

//  vanilla.ts

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

// 定义了一个状态存储的接口,包括获取状态、设置状态、订阅状态变化和销毁状态存储
export interface StoreApi<T> {
  getState: () => T
  setState: SetStateInternal<T>
  // 参数是监听状态值变化函数,返回取消订阅的函数
  subscribe: (listener: (state: T, prevState: T) => void) => () => void
  destroy: () => void
}

// 定义了一个创建状态的函数类型,这个函数接收设置状态、获取状态和状态存储对象作为参数
export type StateCreator<T> = (
  setState: StoreApi<T>['setState'],
  getState: StoreApi<T>['getState'],
  store: StoreApi<T>
) => T

// 创建状态存储的函数
type CreateStore = {
  <T>(createState: StateCreator<T>): StoreApi<T>
  <T>(): (createState: StateCreator<T>) => StoreApi<T>
}

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

type CreateStoreImpl = <T>(createImpl: StateCreator<T>) => StoreApi<T>

export 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 getState: StoreApi<TState>['getState'] = () => state

  // 设置状态的方法,接收一个新的状态或者一个返回新状态的函数作为参数
  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial

    // 如果新的状态和旧的状态不同,那么更新状态,并且通知所有的监听器
    if (!Object.is(state, nextState)) {
      const prevState = state
      state =
        replace ?? typeof nextState !== 'object'
          ? (nextState as TState)
          : Object.assign({}, state, nextState)

      listeners.forEach((listener) => listener(state, prevState))
    }
  }
  //   订阅状态变化
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    return () => listeners.delete(listener)
  }
  //   销毁状态存储
  const destroy: StoreApi<TState>['destroy'] = () => {
    listeners.clear()
  }

  const api = {
    getState,
    setState,
    subscribe,
    destroy
  }
  state = createState(setState, getState, api)

  return api
}

这段代码的主要功能是创建一个状态管理库,可以在React中使用。
通过使用use-sync-external-store库来同步外部的状态,可以在组件之间共享状态。
对应的地址:github.com/facebook/re…

// react.ts
import { StateCreator, StoreApi, createStore } from './vanilla'
import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector'
const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports

// 定义类型,用于提取StoreApi的状态
type ExtractState<S> = S extends { getState: () => infer T } ? T : never

// 定义只读的StoreApi类型
type ReadonlyStoreApi<T> = Pick<StoreApi<T>, 'getState' | 'subscribe'>

// 定义一个可以与React一起使用的StoreApi类型
type WithReact<S extends ReadonlyStoreApi<unknown>> = S & {
  getServerState?: () => ExtractState<S>
}

// 定义一个绑定了Store的Hook的类型
export type UseBoundStore<S extends WithReact<ReadonlyStoreApi<unknown>>> = {
  (): ExtractState<S>
  <U>(
    selector: (state: ExtractState<S>) => U,
    equals?: (a: U, b: U) => boolean
  ): U
} & S

// 定义创建Store的类型
type Create = {
  <T>(createState: StateCreator<T>): UseBoundStore<StoreApi<T>>
  <T>(): (createState: StateCreator<T>) => UseBoundStore<StoreApi<T>>
}

// 创建Store的函数
export const create = function (createState) {
  return createState ? createStoreImpl(createState) : createStoreImpl
} as Create

// 实现创建Store的函数
function createStoreImpl(createState: StateCreator<T>) {
  const api =
    typeof createState === 'function' ? createStore(createState) : createState

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

  return useBoundStore
}

// 使用Store的Hook
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
  )
  return slice
}

至此,我们已经完成了zustand的核心功能实现。借助了react提供的hooks,我们可以用少量的代码实现发布订阅并更新到界面上的功能。