Zustand源码解读(更新中)

底层实现-store对象

底层-类型系统篇

基本类型

在研究store对象如何实现之前 我们先聚焦于类型系统

store类型 也就是api这个对象的类型

export interface StoreApi<T> {
  setState: SetStateInternal<T>
  getState: () => T
  getInitialState: () => T
  subscribe: (listener: (state: T, prevState: T) => void) => () => void
}

api对象上面的四个方法,均定义在这个接口里面。

这个接口是一个泛型接口 接受一个泛型T,这个T代表我们数据对象的类型,也就是creator函数返回的那个对象的类型。这个类型可以是我们手动指定的,但是在数据较为简单清晰的时候,让ts类型系统执行推断也可以。

但是ts类型系统的推断是有局限性的,就例如我们声明一个对象,上面有一个a属性 值为1,实际上这个a的类型可以是一个字面量类型 也可以是number类型,如果没有其他信息那么类型推断系统将会把他推断成number。

通过上面的接口我们可以直观的看到getState getInitailState这两个函数的返回值,以及subscribe应该是一个接受一个函数作为参数 并且返回一个函数的函数 ,(接受一个订阅者函数返回取消订阅的函数 )

从这个store类型 我们就可以看出zustand发布系统设计的一角:

通过subscribe注册监听者,并且返回取消监听的函数

通过setState进行状态更新 并且在更新之后触发监听者。

而setState则是使用了另一个类型 让我们来看看这个类型怎样定义:

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

不是很掌握ts的话 可能会对这里的语法有些误解 看起来是一个怪怪的对象类型

但是这里定义的是函数重载类型:先声明一个对象类型 再取出_对应的函数类型 最终得到的是函数重载类型。

你可能会疑问为什么不直接定义函数重载 就比如:

// 定义一个函数重载类型:支持两种调用方式
type Add = {
  (a: number, b: number): number;       // 重载1:两个数字相加
  (a: string, b: string): string;       // 重载2:两个字符串拼接
};

事实上 这种写法一般用于: 函数参数的类型 依赖于 函数自身的返回值类型的时候

这是因为 如果我们这样写:

type UpdateFn = (
   updater: State | ((prev: State) => State), 
   callback?: (newState: ReturnType<UpdateFn>) => void // ❌ 错误:循环引用
 ) => State;

这个UpdateFn 他的callback参数的类型 依赖于他的返回值。

但是这个时候我们通过returnType获取他的返回值 由于这个函数的定义还没有完成 所以这个返回值类型是不可知的 就会报错

这是因为,TypeScript 无法解析 “一个正在定义的类型的返回值”
这个时候我们就可以通过在外层定义一个对象来解决:

 type UpdateFnWrapper = {
   _: (
     updater: State | ((prev: State) => State),
     // 2. 回调参数类型:引用对象属性 _ 的返回值类型(而非直接引用 UpdateFn)
     callback?: (newState: ReturnType<UpdateFnWrapper['_']>) => void
   ) => State;
 };

为什么这样就可以?因为ts允许在对象属性里面使用对象某个键的属性 及时这个属性没有完成定义 只要不是死循环就行。

说回SetStateInternal这个类型 我们知道了他是定义了两个函数重载 对于第一个 接受的replase是false 代表合并类型更新 第二个代表替换类型更新

第一个接受的partial参数的类型由于是合并类型,可以是完整的store 也可以是部分store,也可以是函数式更新

第二个函数由于是替换更新 所以必须是完整的store

说起这个对象类型 我们知道ts经常会报错:某个对象上面不能具有某种属性

就比如:

interface Iobj{
  a:any
}
b:Iobj = {
  a:1
}
b.c = 2//报错!

在ts里面 能否进行赋值的判据其实有很多 但是总的来说 这个规律就是:精确可以赋值给模糊,所以按照这个规则来说,我们给这个b对象新增属性 实际上在把它变得更加精确

这不应该报错 但事实并不是这样

这其实是ts里面一个特殊的规则,用于避免我们传递一些不必要的属性,但是这个属性只有在直接赋值的时候有效 所以如果我们这样写:

b:Iobj = {
  a:1
}
b = {
  a=1
  b=1
}

就不会报错。

通过这样的类型定义 允许我们使用多种方式进行更新的同时保证类型安全

工具类型

接下来我们来看几个定义的工具类型:

//用于推断getState的返回值类型
export type ExtractState<S> = S extends { getState: () => infer T } ? T : never
//获取T这个对象上面K属性的类型 如果不存在 使用传入的F类型 
type Get<T, K, F> = K extends keyof T ? T[K] : F

我们知道zustand的一大特点是:能够通过中间件进行功能增强 这个增强 指的是:我们先把set get api传递给中间件 中间件对他们进行改写,再把经过改写的版本传递给creator函数。

对于get set函数 经过改写之后他们的函数签名也不会改变,所以在ts层面不需要处理这两个函数。

但是对于api 也就是我们的store对象 类型就需要修改了 ,并且这个修改是根据我们的中间件调用顺序层层递进的 ,其实就和webpack的插件是一样的。

那么这个中间件可能会怎么改写这个api对象?

可能会给他增加一些新的属性方法 也可能会对属性和方法进行改写 ,例如 添加一个add方法,再给这个方法增加一个prevdata参数 ,再把原来number类型的返回值修改成对象类型 。

也就是说 我们的修改 涉及到对store这个对象的属性的修改,以及对这些属性的类型的修改

所以我们每次要做的修改就是:添加以前不具备的属性,进一步修改以前具备的属性,这需要两个工具:

type Cast<T, U> = T extends U ? T : U //判断类型T是否兼容类型U 
type Write<T, U> = Omit<T, keyof U> & U//去掉重复的属性 再更新重复属性

然后我们需要一个工具 接受两个参数 并且根据第二个参数扩充第一个参数 这在源码里面就是StoreMutators

但是你会看到他是一个空类型:

export interface StoreMutators<S, A> {}

这是给开发者留出来的扩展口,因为interface具有合并的特性 你可以不断地向这个接口添加属性 就比如像这样:

export interface StoreMutators<S, A> {
  devtools: Write<Cast<S, object>, { devtools: A }>
  persist: Write<S, { persist: A }>
  immer: Write<Cast<S, object>, object>
  subscribeWithSelector: Write<S, { subscribe: A }>
  // 可以扩展更多中间件
} 

这个interface就像一个switch case 你需要向store类型里面扩充一个中间件 需要添加一个add方法 你就把原来的store类型 以及{add():void}传递给他 再取出正确的属性(就比如中间件persit对应的属性也是persist) 就得到一个正确处理的store

现在我们知道了 怎么把一个中间件类型和store聚合起来 但是我们肯定是不止一个中间件的 那么我们怎么对中间件列表进行处理?

完美兼容ts的关键-Mutate泛型函数

就像我们在js里面写的遍历一样 我们遍历中间件类型 进行扩充即可:

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

先来看一下整体逻辑:

  1. Ms是否为长度不确定的数组?是就返回原始类型
  2. 如果不是 是不是空数组?是就返回原始类型 不是就使用infer关键字推断中间件的标识以及配置
  3. 如果发现数组结构不符合预期 将会返never

其中最关键的逻辑就在于:

 Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>

这里的StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier]就是取出处理好的store类型 并且把还没有处理的数组回传给Mutate。

最终处理了最后一个中间件的时候 Mrs是空数组 回传过去之后 将会走到第二个分支返回最终处理好的store

这就是处理中间件的主要逻辑 现在让我们看看zustand是怎么处理状态初始化函数的类型的:

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 }

这里的Mis(mutator identifiers for store)类型 是一个标识 表明当前应用了哪些中间件

跟踪对 Store API(如 setState getState store 对象) 的改造

记录所有作用于当前 store 的中间件类型,确保中间件对 StoreApi (即 setState getState store 等核心 API)的修改能被 TypeScript 正确识别

当我们自己编写一个中间件的时候 就需要就这个类型进行扩充 例如:

// 1. 定义中间件的标识符(一个独特的类型标记)
type MyMiddlewareIdentifier = { type: 'my-middleware'; name: 'customMethod' }

// 2. 补全中间件的类型定义,显式处理 Mis
const myMiddleware = <T, Mis extends any[], Mos extends any[]>(
  creator: StateCreator<T, Mis, Mos>
): StateCreator<T, [MyMiddlewareIdentifier, ...Mis], Mos> => {
  // 增强 store API,添加 customMethod
  return (set, get, store) => {
    const enhancedStore = {
      ...store,
      customMethod: () => console.log('Custom method called')
    }
    // 将增强后的 store 传递给原始 creator
    return creator(set, get, enhancedStore as any)
  }
}

mos(Mutator Identifiers for State)跟踪对 状态本身(即 StateCreator 函数返回的 U 类型数据) 的改造

每个元素是 [中间件标识, 状态改造逻辑] 的元组,用于描述中间件如何修改状态的类型结构。

就比如 假设我们应用了一个加密中间件:

const encryptMiddleware = (creator) => (set, get, store) => {
  const originalState = creator(set, get, store);
  // 改造状态:加密后返回新结构
  return {
    encrypted: btoa(JSON.stringify(originalState)),
    decrypted: originalState
  };
};

此时,Mos 会记录这个中间件对状态的改造:
Mos = [['encrypt', (state: T) => { encrypted: string; decrypted: T }]]

那么这个mos是怎么得到的? 在最后一行有这样的一句代码:

& { $$storeMutators?: Mos }

这个$$storeMutators是一个函数附加类型 当我们自己编写一个一个中间件的时候 我们就需要对这个附加类型进行扩充,然后ts就可以根据这个附加类型推断出Mos的类型。

T这个泛型参数 代表的是状态的类型 并且是改写之前的状态类型

而U代表的是改写之后的状态类型 也就是 U类型是经过Mos改造之后的T。

所以这个cerator函数最终的返回值是U。

你可能有一个疑问:getState函数的返回值实际上和creator函数是完全一样的 但是getState的返回值不是根据U计算得到的 这难道不会造成getState的返回值是没有经过处理的类型吗?

其实是这样的 当我们观察在源码里面这个createStore的使用方式就会发现:这里的泛型参数都不是指定的 而是使用ts类型推断得到的,

在上面的函数里面 creator函数的返回值将会被推断成U类型 。
我们再看这个泛型参数是怎样申明的:U的默认值为T ,并且T没有其他可以推断其类型的判据了,这会使得T的类型被反向推断成U。

这可能会让你感到很疑惑,那不如直接写T=U呢,反正使用的时候就是推断U 再根据U推断P。

我认为这种误解来自于ts类型,ts类型可以手动指定也可以由类型系统来进行推断,这导致我们有些 颠倒主次。

我认为比较好的解释是: 我们在构思一个类型函数(姑且把泛型类型称为类型函数)的时候,出发点本就不应该是让类型系统来推断,而是就是我们手动传递参数-即使并不这样使用

如果你不知道上面这个函数的使用方式,只看他的参数的话,这个参数表达的含义其实是很明确的,区分了初始的返回值和经过处理的返回值,并且经过处理的返回值默认和初始返回值是相等的 因为获取大部分时间你不需要处理,但是如果你需要的话 也可以进行扩展,但是如果你写T=U 这个参数的含义就让人感到很疑惑了

并且很多人可能对这个中间件的关系有一个误解:认为这个中间件是像链条一样 一个一个处理的 但是并不是的,中间件 包括类型 是全部处理好 再传递给create函数的 所以这个StateCreator类型 也只会调用一次 ,但是似乎很多人都会把这个过程理解成一个层层递进的过程。

柯里化-但是类型

创建store对象 是给调用creteStore函数实现的 而在这个函数里面实际调用的是createStoreImpl

我们先看看createStore函数是怎么定义的:

这个函数按道理来讲应该传递一个creator函数 但是也可以不传递函数 这个时候回返回createStoreImpl

所以这里定义是这样写的:

type CreateStore = {
  //第一个函数重载 接受一个核心数据类型T 中间件列表类型 
  <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>
}

其实这里第二个重载2使用的是柯里化 先传递一个泛型参数 不传递实际参数 然后将会返回一个函数 值的就是散文说到的返回createStoreImpl。

这个函数是这样编写的:

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

你可以看到 在类型定义里面是包含泛型参数的 但是我们实际使用的时候没有传递泛型参数 这是ts的一个小细节:

// 这里的 <T> 是 MyBox 类型的一部分
type MyBox<T> = {
  value: T;
};
// 当我们使用 MyBox 时,必须明确提供 T 的类型
const numberBox: MyBox<number> = { value: 123 };
const stringBox: MyBox<string> = { value: "hello" };
// CreateArray 本身不是泛型,它只是一个类型别名
type CreateArray = {
  // 这里的 <T> 是函数“调用签名”的一部分
  <T>(input: T): T[];
};

// 我们将一个函数断言为 CreateArray 类型
const createArrayFn: CreateArray = (input) => {
  return [input, input];
};

// 当我们“调用”这个函数时,T 才被推断出来
const numberArray = createArrayFn(100);      // T 被推断为 number,函数返回 number[]
const stringArray = createArrayFn("world");  // T 被推断为 string,函数返回 string[]

这个createArray函数 可以接受任何类型的参数 并且这个类型会被T推断出来。

至于CreateStoreImpl 其实就是必须传递参数的createStore:

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

底层-契合框架的发布订阅机制

通过上面的分析我们可以看出 其实两个主要函数就是createStore createStoreImpl

让我们看看createStoreImpl:

const createStoreImpl: CreateStoreImpl = (createState) => {
  //获取返回值类型 也就是状态的ts类型 因为之后函数的参数类型依赖于状态的ts类型
  type TState = ReturnType<typeof createState>
  //这里事件发布订阅机制设计的非常简单 但是这种简单并不是懒得写 而是因为框架内部有逻辑来处理所以zustand这边这样处理就可以了 至于框架做了什么之后再说
  type Listener = (state: TState, prevState: TState) => void
  //闭包维护内部状态
  let state: TState
  //set简化操作
  const listeners: Set<Listener> = new Set()

  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    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 === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

注意看这里:

if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }

当前后的状态不相等的时候 需要执行更新逻辑 这里??是一个空值合并运算符 注意是空值合并 不管真假 是要为真就使用 这里的逻辑是如果显示的传递replace 就直接使用这个标识

当replace为真 直接进行替换 replace是假 直接合并

想不到吧 把ts去掉就只剩五十行代码了

react兼容层

react兼容层-外部状态的矛盾

接下来我们看看react.ts 这是react兼容层 我们先来看具体实现 不看ts类型

首先就是起点 create函数 当然这里的create函数是更高层次的函数:

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

这里用到了之前讲过的定义在参数里面的泛型 这样的泛型不需要我们手动指定 由ts进行推断

createImpl函数:

const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState)

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

  Object.assign(useBoundStore, api)

  return useBoundStore
}

这里的creteStore就是我们刚刚讲到的js层面的函数 并且我们可以看到 我们平时使用的useBoundStore函数是怎样创建的并且把api上的属性传递给他 使得我们可以通过这个函数调用底层方法。

让我们来看看useStore内部是怎么实现的

首先我们明确 我们在使用返回的useBoundStore的时候 会传入一个selector函数 最终得到一个state 这个state根据selector的返回值有所不同 并且这个state改变的时候组件会重新渲染

export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  const slice = React.useSyncExternalStore(
    api.subscribe,
    React.useCallback(() => selector(api.getState()), [api, selector]),
    React.useCallback(() => selector(api.getInitialState()), [api, selector]),
  )
  React.useDebugValue(slice)
  return slice
}

理解这个函数的关键点在于理解react的useSyncExternalStore,

首先 我们为什么需要这个函数?

我们将react内的状态区分成内部状态 外部状态,内部状态指的是使用react内置hook 也就是useState这些hook创建的状态,外部状态指的是从外部系统或者介质里面取出来的状态,比如localStorage。

他们的最关键的区别在于是否直接由react进行管理

对于useState,他创建出来的状态会被保存在一个节点对象上面 并且是放在对象一个数组属性里面的。

但是外部系统里面的变量就不是了,我们每次使用的时候都要取出来。

一个最基本的场景,这个外部系统里面的值更新了,并且需要反映到页面上,这就做不到了,我们顶多每次重新渲染的时候看一下这个值更新没有。

其次,react新版本适用的是fiber架构,我对fiber架构的理解就是:把渲染过程拆分成更多的小单元 也就是非常多的fiber,再拆分成几个不同的阶段,实现更加精细的管理,例如支持并发渲染,能够打断 重新开始

对于打断再重新开始的过程,我们需要保证状态的正确性,对于内部状态这是天然支持的,因为我们会把渲染的阶段拆分成生成组件树阶段 计算虚拟dom阶段 提交阶段 ,假设在生成组件树 计算虚拟dom的阶段发生了中断,我们直接丢弃当前的计算结果就行,但是注意 这个提交阶段是同步的 不可中断的。

但是对于外部状态,要是我们在这两个阶段修改他 又废弃,他是回不到初始值的。

再举一个例子 对于一次渲染过程,他使用了外部状态 并且外部状态值是1,但是这个渲染中断了,重启之后获取到了更新之后的外部状态,因为对于外部状态没有闭包机制来保证稳定性。

所以 基于以上的种种原因 我们需要一个api 安全的订阅外部状态,他要保证在外部状态更新的时候能够触发重新渲染 并且保证外部状态的中断 重启这个过程里面的稳定性

那么我们要怎么解决这个问题?很简单 让外部系统在值更新的时候通知freact,然后react开启新的一轮渲染,而不是在旧的渲染里面使用新数据即可,要保持状态的稳定 我们就使用快照机制,在组件第一次渲染的时候记录快照,之后如果重新渲染就使用之前保存的快照

其实就是谁能准确的知道数据变化,谁就管理数据相关更新。要保证稳定性,就让react创建一份外部状态的副本。

上面的代码里面 这个函数接受了三个参数 我们先看前两个:

 api.subscribe,
 React.useCallback(() => selector(api.getState()), [api, selector]),  

第一个参数是订阅函数,在react内部将会传递一个forceUpdate函数,用于强制组件进行刷新,这样在发布订阅系统发布的时候就会强制组件进行刷新。

第二个函数我们称之为getSnapShot 获取快照 例如我们这里使用store对象里面的count属性 这里这个函数就会返回这个属性

整个函数的运行流程可以这样概括:

  1. **
    **执行 ****getSnapshot:获取外部状态的初始快照(比如从 externalStore.getState() 拿到初始值),作为组件的初始状态。
  2. 执行 ****subscribe ****订阅
    • 传入一个 React 内部的回调函数(我们暂时叫它 forceUpdate)。
    • 外部状态库会保存这个 forceUpdate 回调(比如存入一个监听队列),并返回「取消订阅」的清理函数。
  1. 记录清理函数:React 会记住这个清理函数,在组件卸载或后续不需要订阅时自动调用(避免内存泄漏)。

当外部状态发生变化时(比如调用了 externalStore.setState(...)):

  1. 外部状态库通知订阅者:遍历之前保存的所有 forceUpdate 回调并执行(这一步由外部状态库实现,比如示例中的 listeners.forEach(...))。
  2. React 标记组件需要更新forceUpdate 被调用后,React 会将组件标记为「需要重渲染」,加入更新队列。

进入组件重渲染流程时,useSyncExternalStore 会:

  1. 再次执行 ****getSnapshot:获取最新的状态快照(此时已经是更新后的值)。
  2. 对比快照差异
    • 用浅比较(Object.is)对比「上一次快照」和「当前快照」。
    • 如果不同:使用新快照更新组件,UI 展示最新状态。
    • 如果相同:跳过更新(优化性能)。

在 React 18 的并发渲染模式下(比如使用 Suspense 或过渡动画),组件渲染可能被中断 / 恢复,useSyncExternalStore 会额外做:

  1. 锁定快照版本:在渲染开始时读取一次快照并「锁定」,即使渲染过程中状态又变了,也会用这个锁定的快照完成当前渲染(避免渲染到一半突然换状态,导致 UI 撕裂)。
  2. 渲染完成后校验:如果渲染期间状态确实变了,渲染完成后会立刻触发一次新的更新,保证最终 UI 是最新的。

这样我们就大概明白了这个函数是怎么运行的 那么第三个参数的作用是什么?

第三个参数 我们称之为getServerSnapshot 获取服务端快照,适用于预防在ssr场景下服务端客户端渲染不一致的问题

举个例子,我们的登录状态保存在了locallocalStorage 但是服务端是没有这个东西的 所以只能返回一个没登录 于是服务端发过来的html字符串里面就是没登录 但是客户端接受到之后水合一下,去localStrage里面取出登录状态发现是登录的,于是服务端客户端的状态就不一致了

为了避免这个情况 我们需要在服务端发给我们html字符串之前就把一些必要的数据告诉他

用户访问页面时,服务端可以通过 请求头、Cookie 或 Session 拿到真实的登录状态(比如用户已经登录)。这个状态是 “权威的”,是客户端和服务端的统一基准。

然后getserverSnapShot里面将会使用这些传递的值

至于返回值slice 就是selector函数的返回值 我们称之为状态片段 就是通过快照获取的值。

reaact兼容层-类型系统

接下来我们继续看看ts部分 这部分会比底层的简单一些

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

const identity = <T>(arg: T): T => arg
export function useStore<S extends ReadonlyStoreApi<unknown>>(
  api: S,
): ExtractState<S>

export function useStore<S extends ReadonlyStoreApi<unknown>, U>(
  api: S,
  selector: (state: ExtractState<S>) => U,
): U

这里定义了一个ReadonlyStoreApi 和底层的api对象相比 这个对象限制了只能读取 这是为了避免我们使用有些不再zustand设想的途径去更改状态

如果组件能直接拿到完整的 StoreApi<T> 并调用 setState,可能会绕过 store 内部的逻辑(如中间件、日志、持久化等),导致状态不一致。
ReadonlyStoreApi<T> 相当于 “只读保护”,只允许组件做安全的操作(读状态、订阅变化)。

而useStore函数允许调用的时候传递或者不传递selector 所以这定义了重载

export type UseBoundStore<S extends ReadonlyStoreApi<unknown>> = {
  (): ExtractState<S>
  <U>(selector: (state: ExtractState<S>) => U): 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>>
}

中间件

接下来我们看中间件编写:

persist

这个中间件用于对数据进行持久化 并且支持配置 配置项一共有:

  1. partialize 支持部分状态持久化
  2. storage 指定存储介质
  3. version+migrate 支持状态迁移
  4. merge 进行状态合并 默认浅合并
const useUserStore = create(
  persist(
    // 状态定义:包含用户信息、偏好设置和操作方法
    (set, get) => ({
      // 基本信息(需要持久化)
      userInfo: {
        id: '',
        name: '',
        email: '',
      },
      // 偏好设置(需要持久化)
      preferences: {
        theme: 'light',
        notifications: true,
        // v2 新增字段
        fontSize: 16,
      },
      // 临时状态(不需要持久化)
      isEditing: false,
      tempFormData: {},

      // 设置用户信息
      setUserInfo: (info) => set({ userInfo: { ...get().userInfo, ...info } }),
      // 更新偏好设置
      updatePreferences: (prefs) => set({ 
        preferences: { ...get().preferences, ...prefs } 
      }),
      // 进入编辑模式(临时状态,不持久化)
      startEditing: () => set({ isEditing: true, tempFormData: get().userInfo }),
      // 取消编辑
      cancelEditing: () => set({ isEditing: false, tempFormData: {} }),
    }),
    // 完整的持久化配置
    {
      // 1. 存储的 key 名称(必填)
      name: 'user-storage-v2',

      // 2. 存储介质(默认 localStorage,这里用自定义异步存储)
      storage: createJSONStorage(() => customAsyncStorage),

      // 3. 筛选需要持久化的字段(排除临时状态)
      partialize: (state) => ({
        userInfo: state.userInfo,
        preferences: state.preferences,
        // 明确排除临时字段
      }),

      // 4. 当前状态版本(用于迁移)
      version: 2,

      // 5. 版本迁移函数(处理旧数据结构)
      migrate: (persistedState, version) => {
        // 版本 1 升级到 2:处理新增的 fontSize 字段
        if (version === 1) {
          return {
            ...persistedState,
            preferences: {
              ...persistedState.preferences,
              fontSize: 16, // 新增字段默认值
            },
          };
        }
        // 版本 0 升级到 2:处理更早期的结构
        if (version === 0) {
          return {
            userInfo: persistedState.oldUserInfo || {},
            preferences: {
              theme: persistedState.theme || 'light',
              notifications: persistedState.notifications !== false,
              fontSize: 16,
            },
          };
        }
        return persistedState;
      },

      // 6. 自定义状态合并规则
      merge: (persistedState, currentState) => {
        // 合并逻辑:
        // - 基础字段用存储数据覆盖
        // - 偏好设置采用深层合并(避免覆盖新增字段)
        return {
          ...currentState, // 保留当前状态的临时字段
          ...persistedState, // 覆盖基础持久化字段
          preferences: {
            ...currentState.preferences, // 保留当前默认配置
            ...(persistedState?.preferences || {}), // 合并存储的偏好设置
          },
        };
      },

      // 7. 水合(恢复)前的回调
      onRehydrateStorage: (state) => {
        console.log('开始恢复状态,当前状态:', state);
        // 返回水合完成后的回调
        return (rehydratedState, error) => {
          if (error) {
            console.error('状态恢复失败:', error);
          } else if (rehydratedState) {
            console.log('状态恢复成功:', rehydratedState);
          }
        };
      },

      // 8. 跳过自动水合(默认 false,这里展示用法)
      // skipHydration: false,
    }
  )
);

对于基本功能 也就是支持对部分或者全部的状态进行持久化 其实很简单 存的时候缓存就行了 但是涉及到这么多配置就比较复杂了 我们一步一步看

我们还是先看工具函数:

const toThenable = (fn) => (input) => {
  try {
    const result = fn(input);
    if (result instanceof Promise) {
      return result;
    }
    return {
      then(onFulfilled) {
        return toThenable(onFulfilled)(result);
      },
      catch(_onRejected) {
        return this;
      },
    };
  } catch (e) {
    return {
      then(_onFulfilled) {
        return this;
      },
      catch(onRejected) {
        return toThenable(onRejected)(e);
      },
    };
  }
};

这个函数很清楚:把普通对象转换成thenable对象 通过把同步的调用转换成异步的调用 统一我们处理同步异步的方式

这个函数也是一个柯里化的函数 我们先传递fn 得到一个thenable 再传入值进行调用。 我们来看看函数体是怎么做的:

 const result = fn(input);
    if (result instanceof Promise) {
      return result;
    }
    return {
      then(onFulfilled) {
        return toThenable(onFulfilled)(result);
      },
      catch(_onRejected) {
        return this;
      },
    };

这段很明确:如果是一个promise对象 就可以直接返回 如果不是 返回一个经过包装的对象 这个对象具有then方法 并且当我们传递处理函数 就会再次返回一个thenable

举个例子:

const add = (num)=>num+1
//获取thenable版本
const thenableAdd = thenable(add)
//这个thenableAdd 我们可以传入一个input 之后将会得到一个具有then方法 catch方法的对象 
const obj = thenableAdd(1)
const obj2 = obj.then((res)=>console.log(res))
//...

这个函数可能不是很好理解 可以自己写一下promise的源码理解一下 。

我们再看一下错误处理 这里通过try catch 返回了fn执行过程里面可能出现的同步错误 进入catch块:

catch (e) {
    return {
      then(_onFulfilled) {
        return this;
      },
      catch(onRejected) {
        return toThenable(onRejected)(e);
      },
    };
  }

其实就是模拟promise promise是怎么处理的? 当发生错误的时候 之后的then函数不会再执行 直到找到第一个catch 否则将会抛出错误,这里我们不细讲了,之后再出一个讲promise的。

所以抛出错误的时候 即使再通过then函数传递onFullfilled函数 我们也不应该再执行,而是等待catch函数处理错误的时候 把错误对象传递出去。

所以then函数应该返回对象本身

createJsonParse函数:

export function createJSONStorage(getStorage, options = {}) {
  let storage;
  try {
    storage = getStorage();
  } catch {
    // 服务端渲染等场景下存储不可用时返回 undefined
    return;
  }

  return {
    getItem: (name) => {
      const parse = (str) => {
        if (str === null) return null;
        return JSON.parse(str, options.reviver);
      };

      const str = storage.getItem(name) ?? null;
      if (str instanceof Promise) {
        return str.then(parse);
      }
      return parse(str);
    },
    setItem: (name, newValue) => {
      return storage.setItem(name, JSON.stringify(newValue, options.replacer));
    },
    removeItem: (name) => storage.removeItem(name),
  };
}

这个函数的目的非常简单 简化操作 不需要我们手动序列化 反序列化

并且兼容同步存储 异步存储 ,同步存储比如localStorge 异步比如indexDB

基本工具就是这两个 接下来我们看一看关键逻辑

首先我们明确 一个支持数据持久化的插件 大概运行逻辑是怎么样的?

首先我们配置 对哪些数据需要持久化 持久化的介质,之后每次数据更新的时候,我们都需要更新介质里面的数据

并且用户第一次打开我们的网页 我们还需要检查介质有没有之前持久化的数据 有的话就应用到我们的store里面 这才是持久化最重要的功能

其余的问题 还有如果有的时候我们需要修改持久化数据的数据结构 应该怎样在不同的数据结构里面完成平缓的迁移?毕竟假设有的时候我们的应用更新了 但是用户那里保存的数据还是旧的,如果我们不做迁移 可能会导致数据无法使用 甚至报错。

并且这个中间件 应该是接受配置对象 creator函数 然后返回一个接受set get api参数的函数 我们一步一步来看:

//这个config就是creator函数
const persistImpl = (config, baseOptions) => (set, get, api) => {
  // 默认配置:使用 localStorage 存储,浅合并,全量持久化
  let options = {
    //存储介质 默认就是localStorage
    storage: createJSONStorage(() => localStorage),
    partialize: (state) => state, // 默认持久化全部状态
    //版本 这个版本是数据结构不同 就必须分版本  
    version: 0,
    //状态合并函数 默认就是浅合并
    merge: (persistedState, currentState) => ({
      ...currentState,
      ...(persistedState ?? {}),
    }),
    //跳过水合 水合在编程领域有点像注入的概念 这里的水合就是在程序启动的时候 把介质里面的数据注入react
    skipHydration: false,
    ...baseOptions,
  };

然后我们需要什么? 我们首先想到的应该是一个函数 把状态存储到介质里面 并且返回存储的对象

const setItem = () => {
    const stateToPersist = options.partialize({ ...get() });
    return storage.setItem(options.name, {
      state: stateToPersist,
      version: options.version,
    });
  };

我们需要在每次进行状态更新之后 自动持久化 所以需要重写api上的setState方法 向cretor函数里面传递重写之后的set函数

// ------------------------------
  // 重写 api.setState:状态更新后自动持久化
  // ------------------------------
  const savedSetState = api.setState;
  api.setState = (state, replace) => {
    savedSetState(state, replace);
    return setItem();
  };

  // ------------------------------
  // 执行原始 store 配置,重写 set 方法(触发更新时自动持久化)
  // ------------------------------
  const configResult = config(
    (...args) => {
      set(...args);
      return setItem();
    },
    get,
    api
  );

  // 初始状态为原始配置的返回值
  api.getInitialState = () => configResult;

接下来 我们需要完善水合函数 这也是persist中间件最复杂的一部分

我们先大概构思一下 这个函数要做什么 首先是执行之前 执行前置回调,然后我们第一步应该是从介质里面 读取数据,并且判断取出的数据的版本和现在的版本是不是不一样 如果不一样 我们需要进行迁移 ,然后把完成处理的持久化的数据和现在的数据merge一下,如果持久化数据存在版本变化,还需要更新介质里存储的数据

并且我们之前说到了工具函数toThenable 这里就需要使用它 因为有的介质是同步读取的 有的介质是异步读取的 ,并且借助这个函数 我们可以把这三个过程划分成清晰的三个步骤:

  1. 读取介质数据进行处理
  2. 合并数据
  3. 结束处理

让我们完成这段代码 第一个步骤:

const hydrate = () => {
    if (!storage) return;

    // 标记 hydration 开始,触发监听器
    hasHydrated = false;
    hydrationListeners.forEach((cb) => cb(get() ?? configResult));

    // 执行 hydration 前置回调
    const postRehydrationCallback = options.onRehydrateStorage?.(get() ?? configResult);

    // 从存储读取数据(兼容同步/异步存储)
    return toThenable(storage.getItem.bind(storage))(options.name)
      .then((deserializedStorageValue) => {
        // 1. 处理版本不一致:执行迁移函数 这里的版本 指的是数据结构的改变
        if (deserializedStorageValue) {
          if (
            typeof deserializedStorageValue.version === 'number' &&
            deserializedStorageValue.version !== options.version
          ) {
            if (options.migrate) {
              const migrationResult = options.migrate(
                deserializedStorageValue.state,
                deserializedStorageValue.version
              );
              if (migrationResult instanceof Promise) {
                return migrationResult.then((migratedState) => [true, migratedState]);
              }
              //这个 [true, migrationResult] 数组是一个在 persist 中间件中用于**状态迁移(migration)**的特定返回格式。
              return [true, migrationResult];
            }
            console.error('[zustand persist] 状态版本不匹配,但未提供迁移函数');
          } else {
            // 版本一致:直接返回存储的状态
            return [false, deserializedStorageValue.state];
          }
        }
        // 无存储数据:返回默认值
        return [false, undefined];
      })

然后我们进入第二个阶段 我们将会得到一个数组 数组第一个元素代表是否惊醒迁移 第二个参数代表最终的得到的数据:

.then(([migrated, migratedState]) => {
        // 2. 合并存储状态与当前状态
        stateFromStorage = options.merge(migratedState, get() ?? configResult);
        // 3. 更新 store 状态(强制替换)
        set(stateFromStorage, true);
        // 4. 若执行了迁移,更新存储中的状态
        if (migrated) {
          return setItem();
        }
      })

最终 我们处理后置函数

  .then(() => {
        // 5. 触发 hydration 后置回调
        postRehydrationCallback?.(stateFromStorage, undefined);
        // 6. 更新已 hydrated 状态,触发结束监听器
        stateFromStorage = get();
        hasHydrated = true;
        finishHydrationListeners.forEach((cb) => cb(stateFromStorage));
      })
      .catch((e) => {
        // 7. 捕获错误并通知回调
        postRehydrationCallback?.(undefined, e);
      });
  };

最终 我们通过api对象向外暴露这些方法 以便我们可以在应用中明确持久化状态的现状:是否在进行水合? 更新配置 重置持久化数据:

api.persist = {
    // 更新持久化配置
    setOptions: (newOptions) => {
      options = { ...options, ...newOptions };
      if (newOptions.storage) storage = newOptions.storage;
    },
    // 清空存储的状态
    clearStorage: () => {
      storage?.removeItem(options.name);
    },
    // 获取当前配置
    getOptions: () => options,
    // 手动触发 hydration
    rehydrate: () => hydrate(),
    // 检查是否已完成 hydration
    hasHydrated: () => hasHydrated,
    // 监听 hydration 开始
    onHydrate: (cb) => {
      hydrationListeners.add(cb);
      return () => hydrationListeners.delete(cb);
    },
    // 监听 hydration 结束
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb);
      return () => finishHydrationListeners.delete(cb);
    },
  };

完整的代码可以自行在git上查找