在之前的文章《掌握Zustand:React状态管理的高效之道》中,我们学习了Zustand的基本概念以及它在React中的应用。现在,让我们更进一步,深入探索Zustand的源码,揭开其高效运行背后的秘密!

本文基于Zustand最新代码版本(v5.0.0-rc.2)进行分析,代码仓库地址为:https://github.com/clin211/zustand。

运行环境要求:

  • Node.js >= 12.20.0
  • TypeScript >= 4.5
  • pnpm (推荐版本 v8.12.1,以保证与本文一致)

Zustand源码结构概览

图片

Zustand的源码结构清晰简洁,主要文件和目录如下:

  • src/:核心代码目录
    • middleware/: 中间件目录,包含各种增强功能的中间件:
      • combine.ts: 用于合并多个store的中间件。
      • devtools.ts: 集成Redux DevTools的中间件,方便调试。
      • immer.ts: 支持使用Immer库进行不可变状态管理的中间件。
      • persist.ts: 实现状态持久化的中间件。
      • redux.ts: 模拟Redux工作流程的中间件。
      • subscribeWithSelector.ts: 支持选择性订阅和响应状态变化的中间件。
    • react/: React相关代码目录:
      • shallow.ts: React版本的浅比较函数,用于优化组件渲染。
    • vanilla/: 原生JavaScript版本代码目录:
      • shallow.ts: 原生JavaScript版本的浅比较函数。
    • index.ts: Zustand库的主入口文件,汇总和导出所有核心功能。
    • middleware.ts: 汇总和重新导出所有中间件。
    • react.ts: 实现Zustand与React的绑定,导出React相关的接口和create等方法。
    • shallow.ts: 提供浅比较功能,用于优化性能。
    • traditional.ts: 提供传统的状态管理模式,适用于不使用Hooks的场景。
    • types.d.ts: TypeScript类型声明文件,定义Zustand中使用的类型。
    • vanilla.ts: 实现Zustand的核心功能,不依赖React,使其可在其他环境(如React Native)中使用。
  • docs/:文档目录
  • examples/:示例代码目录
  • tests/:单元测试代码目录
  • package.json:项目配置文件

前端库的核心逻辑通常位于 src 目录,入口文件一般为 index.js/ts,Zustand也不例外。从目录结构可以看出,src 目录下包含 vanillareactmiddleware 三个关键目录。vanilla 指的是不依赖任何框架(如React、Vue等)的纯JavaScript实现,我们可以在许多库的源码中看到类似的 vanilla 文件或文件夹,例如Jotai、Valtio等。它们包含了最基础的JavaScript实现,而我们在React中使用的Zustand版本,则是在这些基础实现之上结合了React Hooks。react 目录负责适配React框架,middleware 则包含各种中间件,用于扩展Zustand的功能。

Zustand源码深度解析

对项目结构有了初步了解后,让我们深入Zustand的源码,逐步揭开它的神秘面纱。首先,我们从入口文件 index.ts 开始。它导出了 vanilla.tsreact.ts 两个文件:

export * from './vanilla.ts';
export * from './react.ts';

由于上一篇文章主要介绍了Zustand在React中的应用,因此本文的源码解读也将从React的视角入手。

VS Code源码阅读小技巧:

  • 操作光标所在文件中的所有代码块:
    • 折叠所有:Command+K+0
    • 展开所有:Command+K+J
  • 操作光标所在代码块内的代码:
    • 折叠:Command+Option+[
    • 展开:Command+Option+]

create函数

首先,我们来看 react.ts 文件。这个文件的代码量并不大,不到70行。将代码折叠后,我们可以清晰地看到导出的函数、方法和类型,其中就包括我们熟悉的 create 方法:

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

图片

这个方法用于创建一个新的Zustand store。当使用 create 方法时,如果传入 createState 参数,则调用 createImpl 方法并将 createState 作为参数传递;否则,直接返回 createImpl 方法本身。createState 是一个回调函数,它包含我们定义的状态和用于更新状态的函数。

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 };

从定义中可以看出,createState 回调函数接收三个参数:setStategetStatestore

createImpl函数

createImpl 函数主要调用了 createStoreuseStore 两个函数:

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

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

  Object.assign(useBoundStore, api);

  return useBoundStore;
};

代码非常简洁,它首先调用 createStore 创建 store 的核心逻辑,然后调用 useStore 创建一个与React组件绑定的hook,最后使用 Object.assignapi (store 的核心操作方法) 合并到 useBoundStore 上,并返回 useBoundStore。接下来,我们分别深入研究 createStoreuseStore 的实现细节。

createStore函数

createImpl 中调用的 createStore 函数位于 vanilla.ts 文件中,源码如下:

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

如果传入 createState 参数,则调用 createStoreImpl 方法并将 createState 作为参数传递;否则,直接返回 createStoreImpl 方法本身。下面我们来看看 createStoreImpl 方法的具体实现:

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) => {
    // 参数如果是函数,那么就是一个函数,这个函数需要接收当前的 state 并返回一个新的 state,否则就赋值
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial;

    // 通过 Object.is 对新state和旧state 进行比较,如果不相等则更新
    if (!Object.is(nextState, state)) {
      const previousState = state;

      // 如果 replace 为 true,则直接赋值,否则合并后一个新对象
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState);

      // 遍历所有的订阅者,通知订阅者
      listeners.forEach((listener) => listener(state, previousState));
    }
  };

  // 获取当前的 state
  const getState: StoreApi<TState>['getState'] = () => state;

  // 获取初始的 state
  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
    initialState;

  // 添加订阅者到订阅者集合里面,同时返回对应销毁函数
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener);
    // Unsubscribe
    return () => listeners.delete(listener);
  };

  // 所有处理函数以对象的形式暴露出去
  const api = { setState, getState, getInitialState, subscribe };
  const initialState = (state = createState(setState, getState, api));
  return api as any;
};

这段代码的核心逻辑是创建一个 store 并提供 setStategetStategetInitialStatesubscribe 等方法来管理状态和订阅状态变化。其中,setState 方法的逻辑可能比较复杂,我们可以将其拆解成以下几个步骤:

  • 如果 partial 参数是一个函数,则将其作为状态更新函数执行,并将当前状态作为参数传递给它,得到新的状态;否则,直接将 partial 作为新的状态。
  • 使用 Object.is 比较新状态和旧状态,如果状态发生变化,则更新状态并通知所有订阅者。
  • 状态更新的逻辑取决于 replace 参数:如果 replacetrue,则直接将新状态赋值给 state;否则,将新状态和旧状态合并成一个新对象,并赋值给 state

useStore函数

接下来,我们来分析 useStore 的实现。

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;

useStore 使用了 TypeScript 的函数重载特性来支持不同的使用方式。第一种方式是最简单的,只传入一个 api 参数,用于获取当前状态。第二种方式传入 apiselector 两个参数,selector 是一个函数,用于从状态中选择需要订阅的部分。

下面我们来看 useStore 的具体实现:

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

这段代码的核心逻辑是使用 React 的 useSyncExternalStore Hook 来订阅外部 store 的状态变化。useSyncExternalStore 接收三个参数:

  • subscribe:用于订阅 store 状态变化的函数。
  • getSnapshot:用于获取当前状态的函数。
  • getServerSnapshot:用于在服务器端渲染时获取初始状态的函数。

在 Zustand 中,subscribe 参数传入的是 api.subscribegetSnapshot 参数传入的是一个函数,该函数调用 api.getState() 获取当前状态,并使用 selector 函数选择需要订阅的部分;getServerSnapshot 参数传入的是一个函数,该函数调用 api.getInitialState() 获取初始状态,并使用 selector 函数选择需要订阅的部分。

最后,useStore 返回的是 slice,即 selector 函数选择的状态片段。

createImpl函数返回值分析

最后,我们再来看看 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;
};

通过上面的分析,我们知道 api 包含了对 store 状态的操作方法,useBoundStore 则是一个用于订阅 store 状态变化的 Hook。createImpl 使用 Object.assignapi 的方法合并到 useBoundStore 上,并返回 useBoundStore。这意味着,我们可以通过 useBoundStore 访问 store 的状态和操作方法,例如:

const useStore = create((set) => ({
  count: 0,
  increment: () => set((state) => { state.count += 1; }),
}));

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

useShallow函数

在上一篇文章中,我们介绍了 useShallow 方法,它可以用于避免不必要的组件重新渲染。当我们需要订阅 store 中的一个计算状态时,推荐的做法是使用 selector 函数,当 selector 函数的返回值发生变化时,组件会重新渲染。useShallow 可以帮助我们优化 selector 函数的性能,它使用浅比较来判断 selector 函数的返回值是否发生变化,只有当返回值发生变化时,才会触发组件重新渲染。

在 React 中,useShallow 的实现是基于 vanilla 版本的 shallow 函数和 React Hook 实现的,它的源码如下:

import React from 'react';
import { shallow } from '../vanilla/shallow.ts';

export function useShallow<S, U>(
  selector: (state: S) => U,
): (state: S) => U {
  // 缓存上一次的值
  const prev = React.useRef<U>();
  return (state) => {
    const next = selector(state);
    // 使用 vanilla 中的 shallow 比较是否要更新
    return shallow(prev.current, next)
      ? (prev.current as U)
      : (prev.current = next);
  };
}

我们再来看看 vanilla 版本的 shallow 函数的实现:

const isIterable = (obj: object): obj is Iterable<unknown> =>
  Symbol.iterator in obj;

// Map 对比
const compareMapLike = (
  iterableA: Iterable<[unknown, unknown]>,
  iterableB: Iterable<[unknown, unknown]>,
) => {
    // 是 Map 就直接返回,如果不是就 new 一个
  const mapA = iterableA instanceof Map ? iterableA : new Map(iterableA);
  const mapB = iterableB instanceof Map ? iterableB : new Map(iterableB);
  
  // 如果 size 不一样就直接返回 false
  if (mapA.size !== mapB.size) return false;
    
  // 逐项对比
  for (const [key, value] of mapA) {
    if (!Object.is(value, mapB.get(key))) {
      return false;
    }
  }
  return true;
};

export function shallow<T>(objA: T, objB: T): boolean {
  // 浅比较是否相等
  if (Object.is(objA, objB)) {
    return true;
  }
  // 做引用类型的判断
  if (
    typeof objA !== 'object' ||
    objA === null ||
    typeof objB !== 'object' ||
    objB === null
  ) {
    return false;
  }

  // 两个对象是否是可迭代的
  if (isIterable(objA) && isIterable(objB)) {
    const iteratorA = objA[Symbol.iterator]();
    const iteratorB = objB[Symbol.iterator]();
    let nextA = iteratorA.next();
    let nextB = iteratorB.next();

    // 两个都是数组且长度都为 2
    if (
      Array.isArray(nextA.value) &&
      Array.isArray(nextB.value) &&
      nextA.value.length === 2 &&
      nextB.value.length === 2
    ) {
      // Map 比较
      return compareMapLike(
        objA as Iterable<[unknown, unknown]>,
        objB as Iterable<[unknown, unknown]>,
      );
    }
    // 逐个比较两个可迭代对象中的元素
    while (!nextA.done && !nextB.done) {
      if (!Object.is(nextA.value, nextB.value)) {
        return false;
      }
      nextA = iteratorA.next();
      nextB = iteratorB.next();
    }
    return !!nextA.done && !!nextB.done;
  }

  // 对比两个对象的键的长度
  const keysA = Object.keys(objA);
  if (keysA.length !== Object.keys(objB).length) {
    return false;
  }

  // 对比两个对象的每个键的值
  for (const keyA of keysA) {
    if (
      !Object.hasOwn(objB, keyA as string) ||
      !Object.is(objA[keyA as keyof T], objB[keyA as keyof T])
    ) {
      return false;
    }
  }
  return true;
}

shallow 函数的实现比较简单,它主要用于浅比较两个对象是否相等。如果两个对象是基本类型或者 null,则直接使用 Object.is 进行比较;如果两个对象都是可迭代对象,则逐个比较它们的元素;否则,比较它们的键和对应的值是否相等。

中间件

combine中间件

combine 中间件用于将多个 store 合并成一个 store,这在管理复杂应用状态时非常有用。它的主要功能是合并两个状态对象,将初始状态对象 initialState 和通过 create 函数生成的状态对象合并成一个新的状态对象,并返回一个新的状态创建函数,该函数会将初始状态和附加状态合并成一个状态对象。

import type { StateCreator, StoreMutatorIdentifier } from '../vanilla.ts';

type Write<T, U> = Omit<T, keyof U> & U;

type Combine = <
  T extends object,
  U extends object,
  Mps extends [StoreMutatorIdentifier, unknown][] = [],
  Mcs extends [StoreMutatorIdentifier, unknown][] = [],
>(
  initialState: T,
  additionalStateCreator: StateCreator<T, Mps, Mcs, U>,
) => StateCreator<Write<T, U>, Mps, Mcs>;

// 返回一个新函数,这个新函数接收可变参数(...a)
export const combine: Combine =
  (initialState, create) =>
  (...a) =>
    Object.assign({}, initialState, (create as any)(...a));

devtools中间件

devtools 中间件用于将 Zustand store 与 Redux DevTools 集成,方便开发者调试状态变化。devtools 中间件的核心实现是对 setState 方法进行包装,并在每次状态更新时向 Redux DevTools 发送状态变化的通知。

图片

这段代码尝试在浏览器环境中获取 Redux DevTools 扩展,如果未安装,则会输出警告信息:

let extensionConnector:
  | (typeof window)['__REDUX_DEVTOOLS_EXTENSION__']
  | false;
try {
  extensionConnector =
    (enabled ?? import.meta.env?.MODE !== 'production') &&
    window.__REDUX_DEVTOOLS_EXTENSION__;
} catch {
  // ignored
}

devtools 中间件的核心逻辑如下:

// 控制是否需要记录状态的变化
let isRecording = true;

// 将 api.setState 替换成一个新的函数并被强制转换为 any 类型。
;(api.setState as any) = ((state, replace, nameOrAction: Action) => {
    
  // 调用原始的 set 函数
  const r = set(state, replace as any);
  
  // 如果不进行任何额外的操作,直接返回 r
  if (!isRecording) return r;
    
  // 根据 nameOrAction 参数来生成一个 action 对象且包含 type 属性
  const action: { type: string } =
    nameOrAction === undefined
      ? { type: anonymousActionType || 'anonymous' }
      : typeof nameOrAction === 'string'
        ? { type: nameOrAction }
        : nameOrAction;
  
  // 如果 store 为 undefined 则通过 connection.send 发送 action 和当前状态并返回 set 后的对象
  if (store === undefined) {
    connection?.send(action, get());
    return r;
  }
  
  // 如果 store 不为 undefined,则发送一个带有 store 信息的 action,将 action.type 更新为 ${store}/${action.type}
  connection?.send(
    {
      ...action,
      type: `${store}/${action.type}`,
    },
    {
      ...getTrackedConnectionState(options.name),
      [store]: api.getState(),
    },
  );
  return r;
}) as NamedSet<S>;

这段代码的作用是在 setState 调用时,记录状态的变化并通过 connection 发送到外部(Redux DevTools)。它提供了对 state 更改的拦截,并根据 isRecordingstore 的状态决定是否发送记录。

通过 actionstore 的组合,它能够为每个状态更改分配唯一的标识符,确保在多 store 的环境下区分不同的状态。

;(
  connection as unknown as {
    // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
    subscribe: (
      listener: (message: Message) => void,
    ) => (() => void) | undefined;
  }
).subscribe((message: any) => {
  // 根据不同的 type 类型进行处理
  switch (message.type) {
    case 'ACTION':
      // 如果 payload 不是字符串就打印错误信息并终止继续执行
      if (typeof message.payload !== 'string') {
        console.error(
          '[zustand devtools middleware] Unsupported action format',
        );
        return;
      }
      // 解析 json 格式的 payload
      return parseJsonThen<{ type: unknown; state?: PartialState }>(
        message.payload,
        (action) => {
          if (action.type === '__setState') {
            if (store === undefined) {
              setStateFromDevtools(action.state as PartialState);
              return;
            }
            // 如果有多个属性则直接打印错误信息
            if (Object.keys(action.state as S).length !== 1) {
              console.error(
                `  
                [zustand devtools middleware] Unsupported __setState action format.  
                When using 'store' option in devtools(), the 'state' should have only one key, which is a value of 'store' that was passed in devtools(),  
                and value of this only key should be a state object. Example: { "type": "__setState", "state": { "abc123Store": { "foo": "bar" } } }  
                `,
              );
            }
            // 从 action.state 中获取 store
            const stateFromDevtools = (action.state as S)[store];
            
            // 如果没有值就直接返回
            if (
              stateFromDevtools === undefined ||
              stateFromDevtools === null
            ) {
              return;
            }
            
            // 如果当前状态与 DevTools 中的状态不同,直接更新应用状态
            if (
              JSON.stringify(api.getState()) !==
              JSON.stringify(stateFromDevtools)
            ) {
              setStateFromDevtools(stateFromDevtools);
            }
            return;
          }

          // 如果 dispatchFromDevtools 不存在,则直接返回
          if (!(api as any).dispatchFromDevtools) return;
          if (typeof (api as any).dispatch !== 'function') return;
            
          // 调用 dispatch 执行 action
          ;(api as any).dispatch(action);
        },
      );

    case 'DISPATCH':
      switch (message.payload.type) {
        // 重置
        case 'RESET':
          setStateFromDevtools(initialState as S);
          if (store === undefined) {
            return connection?.init(api.getState());
          }
          return connection?.init(getTrackedConnectionState(options.name));
        // 提交
        case 'COMMIT':
          if (store === undefined) {
            connection?.init(api.getState());
            return;
          }
          return connection?.init(getTrackedConnectionState(options.name));
        // 回滚
        case 'ROLLBACK':
          return parseJsonThen<S>(message.state, (state) => {
            if (store === undefined) {
              setStateFromDevtools(state);
              connection?.init(api.getState());
              return;
            }
            setStateFromDevtools(state[store] as S);
            connection?.init(getTrackedConnectionState(options.name));
          });
        // 跳转到指定的状态或操作
        case 'JUMP_TO_STATE':
        case 'JUMP_TO_ACTION':
          return parseJsonThen<S>(message.state, (state) => {
            if (store === undefined) {
              setStateFromDevtools(state);
              return;
            }
            if (
              JSON.stringify(api.getState()) !==
              JSON.stringify(state[store])
            ) {
              setStateFromDevtools(state[store] as S);
            }
          });

        // 导入 state
        case 'IMPORT_STATE': {
          const { nextLiftedState } = message.payload;
          const lastComputedState =
            nextLiftedState.computedStates.slice(-1)[0]?.state;
          if (!lastComputedState) return;
          if (store === undefined) {
            setStateFromDevtools(lastComputedState);
          } else {
            setStateFromDevtools(lastComputedState[store]);
          }
          connection?.send(
            null as any, // FIXME no-any
            nextLiftedState,
          );
          return;
        }
        // 暂停或继续记录状态更改
        case 'PAUSE_RECORDING':
          return (isRecording = !isRecording);
      }
      return;
  }
});

这段代码订阅了 Redux DevTools 的消息,并根据接收到的消息类型执行相应的操作,例如更新状态、重置状态、回滚状态等。它实现了 Zustand store 与 Redux DevTools 的双向通信,方便开发者在 Redux DevTools 中查看和操作 Zustand store 的状态。

immer中间件

immer 中间件允许开发者使用 Immer 库来管理 Zustand store 的状态,Immer 可以简化不可变状态的更新操作。

图片

使用 immer 中间件非常简单:

import { produce } from 'immer';

const useStore = create(immer((set) => ({
  count: 0,
  increment: () => set(produce((state) => {
    state.count += 1;
  })),
})));

persist中间件

persist 中间件用于将 Zustand store 的状态持久化到本地存储中(例如 localStorage),并在应用重新加载时恢复状态。

// ...

const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
  // 实现逻辑
};

type PersistImpl = <T>(
  storeInitializer: StateCreator<T, [], []>,
  options: PersistOptions<T, T>,
) => StateCreator<T, [], []>;

// ...
export const persist = persistImpl as unknown as Persist;

persist 中间件的源码较长,这里不完整展示,可以访问 src/middleware/persist.ts 查看完整源码。

persist 中间件的核心是 persistImpl 函数,它通过包装 api.setState 方法,将状态变化同步到本地存储中。当应用重新加载时,它会从本地存储中恢复状态,并更新到 store 中。

persist 中间件的主要流程如下:

  • 初始化持久化选项:将用户传入的 PersistOptions 和默认选项合并。
  • 重写 setState 方法:拦截 Zustand 的 setState 方法,在状态更新后调用 setItem 将更新后的状态持久化到存储中。
  • 初始化状态 hydrate:调用 hydrate 函数,从存储中获取已保存的状态,如果有保存的状态,则合并存储的状态与当前状态。如果版本不匹配且有 migrate 函数,则进行状态迁移。
  • 注册生命周期回调:包括 onHydrateonFinishHydration,分别在状态开始和结束持久化时触发。

persist 中间件还定义了用于存储管理的接口,通过实现这些接口,可以自定义各种存储解决方案(如 localStoragesessionStorage 或其他异步存储)。

export interface StateStorage {
  getItem: (name: string) => string | null | Promise<string | null>;
  setItem: (name: string, value: string) => unknown | Promise<unknown>;
  removeItem: (name: string) => unknown | Promise<unknown>;
}

export type StorageValue<S> = {
  state: S;
  version?: number;
};

PersistOptions 接口定义了持久化的配置项,例如持久化存储位置(默认为 localStorage):

export interface PersistOptions<S, PersistedState = S> {
  /** Name of the storage (must be unique) */
  name: string;
  /** 
   * Use a custom persist storage. 
   * 
   * Combining `createJSONStorage` helps creating a persist storage 
   * with JSON.parse and JSON.stringify. 
   * 
   * @default createJSONStorage(() => localStorage) 
   */
  storage?: PersistStorage<PersistedState> | undefined;
  /** 
   * Filter the persisted value. 
   * 
   * @params state The state's value 
   */
  partialize?: (state: S) => PersistedState;
  /** 
   * A function returning another (optional) function. 
   * The main function will be called before the state rehydration. 
   * The returned function will be called after the state rehydration or when an error occurred. 
   */
  onRehydrateStorage?: (
    state: S,
  ) => ((state?: S, error?: unknown) => void) | void;
  /** 
   * If the stored state's version mismatch the one specified here, the storage will not be used. 
   * This is useful when adding a breaking change to your store. 
   */
  version?: number;
  /** 
   * A function to perform persisted state migration. 
   * This function will be called when persisted state versions mismatch with the one specified here. 
   */
  migrate?: (
    persistedState: unknown,
    version: number,
  ) => PersistedState | Promise<PersistedState>;
  /** 
   * A function to perform custom hydration merges when combining the stored state with the current one. 
   * By default, this function does a shallow merge. 
   */
  merge?: (persistedState: unknown, currentState: S) => S;

  /** 
   * An optional boolean that will prevent the persist middleware from triggering hydration on initialization, 
   * This allows you to call `rehydrate()` at a specific point in your apps rendering life-cycle. 
   * 
   * This is useful in SSR application. 
   * 
   * @default false 
   */
  skipHydration?: boolean;
}

PersistOptions 接口用于配置 Zustand 状态持久化中间件的行为,允许开发者控制状态如何存储、过滤、迁移和合并。核心选项包括存储名称(name),自定义存储方式(storage,如 localStorage),通过 partialize 过滤需要持久化的状态,及在状态恢复(rehydration)前后的回调(onRehydrateStorage)。此外,version 用于控制状态版本,当版本不匹配时可通过 migrate 函数处理旧状态,merge 用于自定义存储状态与当前状态的合并方式,skipHydration 则可防止应用初始化时自动恢复状态。

下面是一个使用 persist 中间件的示例:

import { persist } from 'zustand/middleware';

const useStore = create(
  persist(
    (set) => ({
      user: null,
      login: (user) => set({ user }),
      logout: () => set({ user: null }),
    }),
    {
      name: 'user-storage', // 存储的 key
      partialize: (state) => ({ user: state.user }), // 只持久化 user 字段
      version: 1, // 状态的版本控制
      migrate: (persistedState, version) => {
        // 当版本号不匹配时执行的迁移逻辑
        if (version === 0) {
          return { user: persistedState.user };
        }
        return persistedState;
      },
    }
  )
);

redux中间件

redux 中间件可以让 Zustand 模仿 Redux 的工作流程,提供类似 dispatchreducer 的支持,让开发者能够以 Redux 的方式去管理和更新状态。

type Write<T, U> = Omit<T, keyof U> & U;

// 定义一个 Redux 动作对象,包含一个 type 属性
type Action = { type: string };

type StoreRedux<A> = {
  dispatch: (a: A) => A;
  dispatchFromDevtools: true; // 是否支持从 Redux DevTools 分发 action
};

type ReduxState<A> = {
  dispatch: StoreRedux<A>['dispatch'];
};

type WithRedux<S, A> = Write<S, StoreRedux<A>>;

type Redux = <
  T,
  A extends Action,
  Cms extends [StoreMutatorIdentifier, unknown][] = [],
>(
  reducer: (state: T, action: A) => T,
  initialState: T,
) => StateCreator<Write<T, ReduxState<A>>, Cms, [['zustand/redux', A]]>;