React ---Zustand源码解析

Zustand源码结构：

zustand/
├── docs/                                  # 文档目录
├── examples/                              # 示例代码目录
├── src/                                   # 源代码目录
│   ├── middleware/
│   │   ├── combine.ts    # 实现了将多个store组合在一起的中间件
│   │   ├── devtools.ts   # 实现了与Redux DevTools进行交互的中间件
│   │   ├── immer.ts      # 支持immer库的中间件，用于不可变数据结构
│   │   ├── persist.ts    # 实现了持久化store的中间件
│   │   ├── redux.ts     # 模拟Redux风格的中间件
│   │   └── subscribeWithSelector.ts  # 支持选择订阅和响应的中间件
│   ├── react/
│   │   └── shallow.ts     # 实现 react 相关的浅比较函数，用于优化组件渲染
│   ├── vanilla/
│   │   └── shallow.ts     # 实现原生状态管理版本的浅比较函数
│   ├── context.ts         # 实现 React Context 提供的支持
│   ├── 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
├── tests/                                # 测试代码目录
├── package.json                          # 项目配置文件
├── README.md                             # 项目说明文档
└── ...       

核心概念解析

create方法解析

create 方法位于 src/react.ts 文件中，而它依赖的核心功能主要在 src/vanilla.ts 文件中实现。

功能和作用

create 方法的主要作用是初始化一个 Zustand store。它接受一个状态初始化函数，并返回一个包含状态管理方法的对象。其职责包括：

1. 初始化状态。
2. 提供 getState 和 setState 方法来读取和更新状态。
3. 提供 subscribe 方法来订阅状态变化。
4. 绑定 React 提供 useStore hook。

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

create 方法封装了 Zustand 的 store 创建和状态管理逻辑，使得使用者可以通过简单直接的方式来管理状态。

Store的创建与初始化

在Zustand中，store的创建是通过create方法实现的。createStoreImpl函数用于完成状态管理逻辑的初始化，首先声明一个状态对象(state)、一个监听器集合（listeners）以及几个核心函数(setState、getState、subscribe等)。这些核心函数允许外部组件不仅可以获取当前的状态，还可以更新状态，并在状态更新时接收通知。

通过调用初始化函数createState,用户定义的状态和更新逻辑会被绑定到核心函数上，完成store的配置。

// src/vanilla.ts

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) => {
    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));
    }
  };

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

  const getInitialState: StoreApi<TState>['getInitialState'] = () => initialState;

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

  const destroy: StoreApi<TState>['destroy'] = () => {
    console.warn(
      '[DEPRECATED] The `destroy` method will be unsupported in a future version. Instead use the unsubscribe function returned by subscribe.'
    );
    listeners.clear();
  };

  const api = { setState, getState, getInitialState, subscribe, destroy };
  const initialState = (state = createState(setState, getState, api));
  return api as any;
};

初始状态管理

getInitialState函数提供机制，用于在store创建时获取初始状态。通过在createState调用中，利用用户提供的状态定义实现。确保创建store，获得一致的初始状态。

const getInitialState: StoreApi<TState>['getInitialState'] = () => initialState;

状态订阅与更新机制

订阅的实现

通过subscribe函数，外部组件可以将一个监听器(listener)添加到内部的监听器集合中。每当状态更新时（通过setState），所有的监听器都会被调用，以通知状态的变化。

const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
  listeners.add(listener);
  // 返回取消订阅的函数
  return () => listeners.delete(listener);
};

监听器在注册时会返回一个取消订阅的函数，调用该函数可以从监听器集合中移除监听器。

状态的推送和变更流程

当调用setState函数更新状态时，会根据提供的新状态或更新函数计算出下一个状态。如果新状态与当前状态不同，则会将新状态设置为当前状态，遍历调用所有注册的监听器，传递新状态和前一个状态作为参数。

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

在setState中，可以选择替换整个状态或部分更新。更新操作完成后，遍历所有监听器并调用它们，通知状态发生变化。

React组件订阅

React组件通过useStore hook内部订阅store实现订阅和响应状态变化的。通过useSyncExternalStoreWithSelector实现，该函数用于在状态变化时通知React组件进行重新渲染.

// src/traditional.ts

import ReactExports from 'react';
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 { useDebugValue } = ReactExports;
const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports;

...

export function useStore<TState, StateSlice  api: WithReact<ReadonlyStoreApi<TState>>,
  selector: (state: TState) => StateSlice = identity as any,
  equalityFn?: (a: StateSlice, b: StateSlice) => boolean,
) {
  if (
    import.meta.env?.MODE !== 'production' &&
    equalityFn &&
    !didWarnAboutEqualityFn
  ) {
    console.warn(
      "[DEPRECATED] Use `createWithEqualityFn` instead of `create` or use `useStoreWithEqualityFn` instead of `useStore`. They can be imported from 'zustand/traditional'. https://github.com/pmndrs/zustand/discussions/1937",
    );
    didWarnAboutEqualityFn = true;
  }
  
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getInitialState,
    selector,
    equalityFn,
  );
  
  useDebugValue(slice);
  return slice;
}

• 1:useSyncExternalStoreWithSelector:该函数是React官方提供的useSyncExternalStore hook的增强版本，添加对选择器函数的支持。通过该函数，React组件可以订阅外部store并在状态变化时重新渲染。

• 2:订阅流程：当组件调用useStore hook时，该hook内部会调用useSyncExternalStoreWithSelector,并传入以下几个参数：

  • 2.1:api.subscribe:用于订阅store状态变化的函数。每当状态通过setState发生变化时，会调用所有已订阅的回调函数。
  • 2.2: api.getState:用于获取当前store状态的函数。
  • 2.3:api.getInitialState:用于获取store初始状态的函数
  • 2.4:selector:用于从状态中提取需要的slice的函数
  • 2.5:equalityFn:用于比较两个状态是否相同的函数，如果两个状态相同则不会触发重渲染。

中间件支持

combine

• 功能：将多个store 组合在一起，适用于复杂状态管理。
• combine函数用来将两个状态对象组合成一个状态对象，接受两个参数
  • 初始状态initialState
  • 附加状态创建函数additionalStateCreator. 返回一个新的状态创建函数，这个新的状态创建函数会将初始状态和附加状态合并为一个状态对象。

devtools

• 功能：提供与Redux DevTools进行交互的功能，便于调试状态变化。
• devtools中间件允许开发者通过Redux DevTools工具监控状态变化和调试。

immer

• 功能：支持使用Immer库实现不可变状态更新。
• immer中间件允许使用Immer库管理不可变状态。immer接受一个状态创建函数，并在内部通过produce函数生成不可变的新状态。

示例：

import {create} from 'zustand'
const seStore = create(immer((set) => ({
  count:0
  increment:() => set(produce((state) => {
   state.count += 1;
  }))
})))

不用immer的写法

import {create} from 'zustand'

const useCountStore = create(){
 (set)=>({
  couint:0,
  increment:()=> set((state) =>({count:state.count ++}))
  
 })
}

使用immer的写法

// npm i immer
import {create} from 'zustand'
import {immer} from 'zustand/middleware/immer'
const useCountStore = create()
 immer((set)=>({
  couint:0,
  increment:()=> set((state) =>{state.count ++}
  
 })
)

persist

• 功能：persost中间件用于实现状态持久化。将store的状态保存在本地存储中(如localStorage),并在应用重新加载时恢复状态。

示例：

const useStore = create(persist((set) =>({
  count:0,
  increment:() => set((state) => ({count:state.count + 1 }))
}),
 {
  name:'count-storage',  // 定义存储名称
  getStorage:() => localStorage, //使用localStorage作为存储
 }
))

redux

• 功能：redux中间件的核心在于通过dispatch方法将动作传递给reducer,根据reducer的返回值更新状态。

示例：

import create from 'zustand';
import { redux } from 'zustand/middleware';

// 定义 reducer 函数
const reducer = (state, action) => {
  switch (action.type) {
    case 'INCREMENT':
      return { count: state.count + 1 };
    case 'DECREMENT':
      return { count: state.count - 1 };
    default:
      return state;
  }
};

// 创建 store
const useStore = create(redux(reducer, { count: 0 }));

function Counter() {
  const count = useStore((state) => state.count);
  const dispatch = useStore((state) => state.dispatch);

  return (
    <div>
      <span>{count}</span>
      <button onClick={() => dispatch({ type: 'INCREMENT' })}>Increment</button>
      <button onClick={() => dispatch({ type: 'DECREMENT' })}>Decrement</button>
    </div>
  );
}

subscribeWithSelector

• 功能：subscribeWithSelector中间件允许通过选择器订阅状态变化，便于在特定的状态子集变化时触发回调。

示例：

import create from 'zustand';
import { subscribeWithSelector } from 'zustand/middleware';

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

function Counter() {
  const count = useStore((state) => state.count);
  const increment = useStore((state) => state.increment);

  return (
    <div>
      <span>{count}</span>
      <button onClick={increment}>Increment</button>
    </div>
  );
}

function TextInput() {
  const text = useStore((state) => state.text);
  const setText = useStore((state) => state.setText);

  return (
    <input
      value={text}
      onChange={(e) => setText(e.target.value)}
    />
  );
}

常见问题与解决方案

如何在多个组件间共享Zustand状态

使用zustand创建的store是全局的，可以在任何组件中通过useStore hook访问和更新状态

// store.js
import create from 'zustand';

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

export default useStore;

// ComponentA.js
import React from 'react';
import useStore from './store';

function ComponentA() {
  const count = useStore((state) => state.count);
  return <div>Count: {count}</div>;
}

export default ComponentA;

// ComponentB.js
import React from 'react';
import useStore from './store';

function ComponentB() {
  const increment = useStore((state) => state.increment);
  return <button onClick={increment}>Increment</button>;
}

export default ComponentB;

如何在应用刷新后保持zustand的状态？

使用zustand/middleware的persist中间件将状态存储在本地存储，如localStorage.

import create from 'zustand';
import { persist } from 'zustand/middleware';

const useStore = create(persist(
  (set) => ({
    count: 0,
    increment: () => set((state) => ({ count: state.count + 1 })),
  }), 
  {
    name: 'count-storage',  // 存储的名字
    getStorage: () => localStorage,  // 存储位置
  }
));

export default useStore;

如何在zustand中使用异步操作？

可以在状态更新函数内部处理异步操作，如API请求

import create from 'zustand';

const useStore = create((set) => ({
  items: [],
  fetchItems: async () => {
    const response = await fetch('https://api.example.com/items');
    const data = await response.json();
    set({ items: data });
  },
}));

export default useStore;

如何在数组或对象中更新特定的项？

通过setState函数和spread操作符实现部分更新

import create from 'zustand';

const useStore = create((set) => ({
  users: [{ id: 1, name: 'Alice' }, { id: 2, name: 'Bob' }],
  updateUser: (id, newName) =>
    set((state) => ({
      users: state.users.map(user => user.id === id ? { ...user, name: newName } : user)
    })),
}));

export default useStore;