手写 zustand

目标功能:

  1. create(initializer): 创建 store 的核心函数。
  2. initializer(set, get, api): 用户定义的初始化函数,接收 set, getapi 对象。
  3. set(partialState | updateFn, replace?): 更新状态,支持部分更新、函数式更新和完全替换。
  4. get(): 在 actions 内部同步获取当前状态。
  5. subscribe(listener, selector?, equalityFn?): 订阅状态变化,支持选择器和自定义比较函数。
  6. useStore(selector?, equalityFn?): React Hook,用于在组件中订阅状态,并进行优化渲染。
  7. api.getState(): 获取当前状态(在 Hook 外部)。
  8. api.setState(): 更新状态(在 Hook 外部)。
  9. api.subscribe(): 订阅状态变化(在 Hook 外部)。
  10. api.destroy(): 清理 store(移除所有监听器)。
  11. 基础中间件支持(通过包装 initializer 实现)。
  12. 使用 useSyncExternalStore (如果可用) 或 useState/useEffect 实现 Hook,确保与 React 并发特性兼容。

项目结构 (概念性):

perl 复制代码
my-zustand-like/
├── src/
│   ├── index.js         # 主入口,导出 create 函数
│   ├── createStore.js   # 核心的 store 创建逻辑 (set, get, subscribe, state management)
│   ├── useStore.js      # React Hook 的实现
│   └── middleware/      # 中间件示例 (可选)
│       └── logger.js
└── package.json

代码实现 (src/index.js 或合并在一个文件中)

js 复制代码
// my-zustand-like.js

// -----------------------------------------------------------------------------
// 0. 依赖项和辅助工具 (模拟 React Hooks 和 useSyncExternalStore)
// -----------------------------------------------------------------------------

// 这是一个非常简化的 React Hooks 模拟,仅用于演示目的
// 实际应用中我们会直接使用 React 提供的 Hooks
const React = (() => {
  let R;
  try {
    R = require('react'); // 尝试加载真实的 React
    console.log("Using actual React library.");
  } catch (e) {
    console.warn("Actual React library not found. Using basic simulation.");
    // 如果没有安装 React,提供一个极简模拟
    let stateId = 0;
    const states = [];
    let effectId = 0;
    const effects = []; // { deps: [], cleanup: null | (() => void) }

    R = {
      useState: (initialValue) => {
        const id = stateId++;
        states[id] = states[id] ?? initialValue;
        const setState = (newValue) => {
          const oldValue = states[id];
          if (!Object.is(oldValue, newValue)) {
            console.log(`[React Simulation] setState for ID ${id}:`, oldValue, '->', newValue);
            states[id] = newValue;
            // 在真实 React 中,这里会触发组件重新渲染
            // 在模拟中,我们无法真正触发渲染,但可以打印日志
            console.log(`[React Simulation] Pretending to re-render component using state ID ${id}`);
          }
        };
        return [states[id], setState];
      },
      useEffect: (callback, deps) => {
        const id = effectId++;
        const oldEffect = effects[id];
        let needsRun = !oldEffect; // 首次运行
        if (oldEffect && deps) {
          // 检查依赖项是否变化
          if (!deps.every((dep, i) => Object.is(dep, oldEffect.deps[i]))) {
            needsRun = true;
            if (typeof oldEffect.cleanup === 'function') {
              console.log(`[React Simulation] Running cleanup for effect ID ${id}`);
              oldEffect.cleanup();
            }
          }
        } else if (oldEffect && !deps) {
            // 没有依赖项,每次都运行
            needsRun = true;
             if (typeof oldEffect.cleanup === 'function') {
              console.log(`[React Simulation] Running cleanup for effect ID ${id}`);
              oldEffect.cleanup();
            }
        }

        if (needsRun) {
          console.log(`[React Simulation] Running effect ID ${id}`);
          const cleanup = callback();
          effects[id] = { deps, cleanup };
        } else {
             console.log(`[React Simulation] Skipping effect ID ${id} (deps unchanged)`);
        }
      },
      useRef: (initialValue) => {
        // 简化 useRef 模拟
        const [ref] = R.useState({ current: initialValue });
        return ref;
      },
      // 尝试使用 React 18 的 useSyncExternalStore,如果不可用则提供 shim
      useSyncExternalStore:
        typeof require === 'function' && (() => { // 仅在 Node.js 环境或打包器环境中尝试 require
          try {
            return require('react').useSyncExternalStore;
          } catch { return undefined; }
        })() ??
        // useSyncExternalStore 的基本 shim 实现 (基于 useState/useEffect)
        // 注意:这个 shim 不完全符合 React 并发模式的所有保证,但能模拟基本行为
        (subscribe, getSnapshot, getServerSnapshot) => {
          const value = getSnapshot();
          const [{ inst }, forceUpdate] = R.useState({ inst: { value, getSnapshot } });

          R.useEffect(() => {
            // 每次渲染后检查快照是否已更改,如果更改则强制更新
            if (!Object.is(inst.value, inst.getSnapshot())) {
                console.log('[useSyncExternalStore Shim] Snapshot changed during render, forcing update.');
                inst.value = inst.getSnapshot(); // 更新内部值
                forceUpdate({ inst });
            }

            // 订阅外部 store
            const handleStoreChange = () => {
              const nextValue = inst.getSnapshot();
              if (!Object.is(inst.value, nextValue)) {
                 console.log('[useSyncExternalStore Shim] Store changed, forcing update.');
                 inst.value = nextValue; // 更新内部值
                 forceUpdate({ inst });
              } else {
                  console.log('[useSyncExternalStore Shim] Store changed, but snapshot is the same. Skipping update.');
              }
            };

            console.log('[useSyncExternalStore Shim] Subscribing to external store.');
            const unsubscribe = subscribe(handleStoreChange);

            // 组件卸载时取消订阅
            return () => {
                console.log('[useSyncExternalStore Shim] Unsubscribing from external store.');
                unsubscribe();
            };
          }, [subscribe]); // 依赖项是 subscribe 函数引用

          // 更新 getSnapshot 函数引用 (如果它变化了)
          inst.getSnapshot = getSnapshot;
          inst.value = value; // 确保 value 是最新的

          return value;
        },
    };
  }
  return R;
})();

// 默认的浅比较函数 (类似 zustand/shallow)
function shallowCompare(objA, objB) {
  if (Object.is(objA, objB)) {
    return true;
  }
  // 必须都是对象且不为 null
  if (typeof objA !== 'object' || objA === null || typeof objB !== 'object' || objB === null) {
    return false;
  }

  const keysA = Object.keys(objA);
  const keysB = Object.keys(objB);

  if (keysA.length !== keysB.length) {
    return false;
  }

  for (let i = 0; i < keysA.length; i++) {
    const key = keysA[i];
    // 检查 key 是否存在于 B 中,并且值是否浅相等 (使用 Object.is)
    if (!Object.prototype.hasOwnProperty.call(objB, key) || !Object.is(objA[key], objB[key])) {
      return false;
    }
  }

  return true;
}


// -----------------------------------------------------------------------------
// 1. Store 核心实现 (createStoreApi)
//    这部分逻辑不依赖 React,是 Zustand 的纯 JS 核心
// -----------------------------------------------------------------------------

/**
 * 创建 Store 的核心 API 和状态管理逻辑。
 * @template TState The type of the state.
 * @param {function(
 *   function(TState | function(TState): TState, boolean?): void,
 *   function(): TState,
 *   StoreApi<TState>
 * ): TState} createState - 用户提供的初始化函数,返回初始状态和 actions。
 * @returns {StoreApi<TState>} Store 的 API 对象。
 */
function createStoreApi(createState) {
  // --- 内部状态和监听器 ---
  /** @type {TState} 保存当前 store 的状态 */
  let state;
  /** @type {Set<function(TState, TState): void>} 存储所有监听 state 变化的函数 */
  const listeners = new Set();

  // --- 核心方法 ---

  /**
   * 更新 Store 状态并通知监听器。
   * @param {TState | function(TState): TState} partial - 可以是部分状态对象,也可以是一个接收当前状态并返回部分状态的函数。
   * @param {boolean} [replace=false] - 如果为 true,则完全替换整个 state,而不是合并。
   * @param {string | object} [actionType] - (可选) 用于调试或中间件的 action 标识。
   */
  const setState = (partial, replace = false, actionType) => {
    // 保存更新前的状态,用于通知监听器
    const previousState = state;

    // 计算下一个状态片段
    // 如果 partial 是函数,调用它获取状态片段;否则直接使用 partial 对象
    const nextStateFragment = typeof partial === 'function' ? partial(state) : partial;

    // 更新状态
    // 如果 replace 为 true,或者 nextStateFragment 和当前 state 是同一个对象 (函数式更新可能返回原对象)
    // 则进行状态合并;否则直接替换。
    // 注意:合并是浅合并。
    if (Object.is(state, nextStateFragment) && !replace) {
        console.warn("[MyZustand] setState called with the same state object reference and replace=false. No actual state change occurred.");
        return; // 如果状态引用没变且不是替换,则不执行任何操作
    }

    // 执行状态更新:替换或合并
    state = replace ? nextStateFragment : Object.assign({}, state, nextStateFragment);

    // --- 通知监听器 ---
    // 使用 Promise.resolve().then() 或 queueMicrotask() 将通知推迟到微任务队列
    // 这有助于批量处理更新,并确保在同步代码执行完毕后才通知
    // 模拟 React 的批处理行为
    Promise.resolve().then(() => {
        if (!Object.is(previousState, state)) { // 只有状态真正改变时才通知
            listeners.forEach((listener) => {
                // 将新状态和旧状态都传递给监听器,方便进行比较
                try {
                    listener(state, previousState);
                } catch (error) {
                    console.error("[MyZustand] Error in listener:", error);
                }
            });
             console.log(`[MyZustand] Notified ${listeners.size} listeners. Action:`, actionType || '(Not specified)');
             console.log("[MyZustand] Previous State:", previousState);
             console.log("[MyZustand] New State:", state);
        } else {
             console.log("[MyZustand] State reference changed, but content is identical after update. Skipping listener notification.");
        }
    }).catch(err => console.error("[MyZustand] Error in microtask scheduling:", err));
  };

  /**
   * 同步获取当前 Store 的状态。
   * 主要用于在 actions 内部获取最新状态。
   * @returns {TState} 当前状态。
   */
  const getState = () => state;

  /**
   * 订阅 Store 的状态变化。
   * @param {function(TState, TState): void} listener - 状态变化时调用的回调函数,接收新状态和旧状态。
   * @returns {function(): void} 取消订阅函数。
   */
  const subscribe = (listener) => {
    // 将监听器添加到 Set 中
    listeners.add(listener);
    console.log(`[MyZustand] Listener added. Total listeners: ${listeners.size}`);
    // 返回一个取消订阅的函数
    return () => {
      listeners.delete(listener);
      console.log(`[MyZustand] Listener removed. Total listeners: ${listeners.size}`);
    };
  };

  /**
   * 销毁 Store,移除所有监听器。
   */
  const destroy = () => {
    const count = listeners.size;
    listeners.clear();
    console.log(`[MyZustand] Store destroyed. Removed ${count} listeners.`);
  };

  // --- API 对象 ---
  // 这个对象包含了 store 的核心方法,可以传递给 createState,也可以附加到 useStore Hook 上
  /** @type {StoreApi<TState>} */
  const api = { setState, getState, subscribe, destroy };

  // --- 初始化 State ---
  // 调用用户传入的 createState 函数,传入 set, get 和 api 对象
  // createState 函数负责返回初始状态(包含数据和 actions)
  state = createState(setState, getState, api);
  console.log("[MyZustand] Store initialized with state:", state);

  // 返回 API 对象
  return api;
}


// -----------------------------------------------------------------------------
// 2. React Hook 实现 (useStore)
//    这部分将 Store 核心 API 与 React 组件连接起来
// -----------------------------------------------------------------------------

/**
 * 创建与特定 Store API 关联的 React Hook。
 * @template TState
 * @param {StoreApi<TState>} api - 由 createStoreApi 创建的 Store API 对象。
 * @returns {UseBoundStore<TState>} 用于在 React 组件中访问和订阅 Store 的 Hook。
 */
function createUseStoreHook(api) {
  // 从 API 对象中解构出核心方法
  const { getState, subscribe: subscribeToStore } = api;

  /**
   * React Hook,用于订阅 Store 状态的选定部分。
   * @template TSlice
   * @param {function(TState): TSlice} [selector = state => state] - 选择器函数,用于从完整状态中提取组件关心的部分。默认为返回整个状态。
   * @param {function(TSlice, TSlice): boolean} [equalityFn = Object.is] - 比较函数,用于判断选择的状态部分是否发生变化。默认为严格相等比较 (Object.is)。可以使用 shallowCompare。
   * @returns {TSlice} 返回选定的状态部分。
   */
  function useStore(selector = (s) => s, equalityFn = Object.is) {
    // --- 使用 useSyncExternalStore (推荐方式) ---
    // 它能更好地处理并发渲染和外部状态同步

    // 1. `subscribe` 函数:
    //    这个函数会被 useSyncExternalStore 调用,用于订阅 store 的变化。
    //    它接收一个由 React 提供的回调 `onStoreChange`。
    //    我们需要在 store 真正变化 *并且* selector 选择的值也变化时,才调用 `onStoreChange`。
    const subscribe = React.useCallback((onStoreChange) => {
        console.log(`[useStore Hook (${selector.toString()})] Subscribing via useSyncExternalStore.`);

        // 内部监听器,当 store 变化时被调用
        const listener = (newState, previousState) => {
            // 使用当前的 selector 和 equalityFn 来比较新旧 slice
            try {
                const newSlice = selector(newState);
                const oldSlice = selector(previousState); // 从旧状态计算旧 slice

                if (!equalityFn(newSlice, oldSlice)) {
                    console.log(`[useStore Hook (${selector.toString()})] Selected state changed, notifying React.`);
                    // 只有当选定的 slice 变化时,才通知 React
                    onStoreChange();
                } else {
                    console.log(`[useStore Hook (${selector.toString()})] Store changed, but selected state is the same. Skipping React notification.`);
                }
            } catch (error) {
                console.error(`[useStore Hook (${selector.toString()})] Error in selector or equalityFn during state update:`, error);
                // 发生错误时,可能需要通知 React 以避免状态不一致,但这取决于具体策略
                onStoreChange(); // 谨慎起见,通知 React
            }
        };
        // 调用 store 核心的 subscribe 方法
        const unsubscribe = subscribeToStore(listener);
        return unsubscribe; // 返回取消订阅函数
    }, [subscribeToStore, selector, equalityFn]); // 依赖项包含 selector 和 equalityFn,确保它们变化时能重新订阅

    // 2. `getSnapshot` 函数:
    //    返回当前状态的快照 (应用了 selector)。
    const getSnapshot = () => {
        try {
            return selector(getState());
        } catch (error) {
            console.error(`[useStore Hook (${selector.toString()})] Error in selector during getSnapshot:`, error);
            // 根据错误处理策略,可能返回一个默认值或上一个有效值
            // 这里简单地重新抛出错误
            throw error;
        }
    };

    // 3. `getServerSnapshot` 函数 (可选):
    //    在服务器端渲染 (SSR) 时使用。对于客户端 store,通常与 getSnapshot 相同。
    const getServerSnapshot = getSnapshot;


    // 调用 useSyncExternalStore
    const selectedState = React.useSyncExternalStore(
      subscribe,
      getSnapshot,
      getServerSnapshot
    );

    // --- (备选方案) 使用 useState 和 useEffect 的实现 ---
    // 这个版本更容易理解基本原理,但在并发模式下可能不如 useSyncExternalStore 健壮
    /*
    const [slice, setSlice] = React.useState(() => selector(getState()));
    const stateRef = React.useRef(getState()); // 保存 state 引用
    const selectorRef = React.useRef(selector);
    const equalityFnRef = React.useRef(equalityFn);
    const sliceRef = React.useRef(slice); // 保存 slice 引用

    // 更新 refs (确保 effect 中拿到最新的)
    selectorRef.current = selector;
    equalityFnRef.current = equalityFn;
    stateRef.current = getState(); // 每次渲染都获取最新 state
    sliceRef.current = slice; // 每次渲染都更新 slice

    React.useEffect(() => {
      console.log(`[useStore Hook (${selectorRef.current.toString()})] useEffect setup: Subscribing.`);
      const listener = (currentState, previousState) => {
        const currentSelector = selectorRef.current;
        const currentEqualityFn = equalityFnRef.current;
        const previousSlice = sliceRef.current; // 获取上次渲染时的 slice
        try {
            const nextSlice = currentSelector(currentState); // 计算新 slice
            if (!currentEqualityFn(previousSlice, nextSlice)) {
              console.log(`[useStore Hook (${currentSelector.toString()})] Slice changed, updating component state.`);
              sliceRef.current = nextSlice; // 更新 ref
              setSlice(nextSlice); // 触发组件重渲染
            } else {
                 console.log(`[useStore Hook (${currentSelector.toString()})] Slice unchanged, skipping update.`);
            }
        } catch (error) {
             console.error(`[useStore Hook (${currentSelector.toString()})] Error during state update comparison:`, error);
             // 可能需要设置错误状态或采取其他措施
        }
      };

      const unsubscribe = subscribeToStore(listener);
      console.log(`[useStore Hook (${selectorRef.current.toString()})] Subscribed.`);

      // 组件卸载时取消订阅
      return () => {
        console.log(`[useStore Hook (${selectorRef.current.toString()})] useEffect cleanup: Unsubscribing.`);
        unsubscribe();
      };
    }, [subscribeToStore]); // 依赖项只有 subscribeToStore (其引用是稳定的)
    // 注意:如果 selector 或 equalityFn 变化,理论上应该重新运行 effect 来确保比较逻辑正确,
    // 但在简单实现中,我们依赖 ref 来获取最新函数。useSyncExternalStore 处理这个更优雅。

    // 每次渲染时,重新计算 slice 并与当前 state 中的 slice 比较
    // 这有助于处理 selector 或 state 在渲染期间同步变化的情况
    let currentSliceFromState;
    try {
        currentSliceFromState = selector(getState());
    } catch (error) {
        console.error(`[useStore Hook (${selector.toString()})] Error in selector during render:`, error);
        throw error;
    }
    if (!equalityFn(slice, currentSliceFromState)) {
        console.warn(`[useStore Hook (${selector.toString()})] Detected change during render. Updating slice.`);
        // 如果渲染期间计算的 slice 与 state 中的不同,立即更新
        setSlice(currentSliceFromState);
        return currentSliceFromState; // 返回最新的
    }

    return slice; // 返回当前组件 state 中的 slice
    */

    // 返回通过 useSyncExternalStore 获取的、保证同步的最新选定状态
    return selectedState;
  }

  // 将 store 的核心 API 方法附加到 useStore Hook 函数本身
  // 这样可以通过 `useCustomStore.getState()` 等方式在组件外部访问 store
  Object.assign(useStore, api);

  // 类型标注 (JSDoc for better DX)
  /**
   * @typedef {object} StoreApi
   * @property {function(TState | function(TState): TState, boolean?): void} setState
   * @property {function(): TState} getState
   * @property {function(function(TState, TState): void): function(): void} subscribe
   * @property {function(): void} destroy
   */

  /**
   * @typedef {function(function(TState): TSlice, function(TSlice, TSlice): boolean?): TSlice & StoreApi<TState>} UseBoundStore
   * @template TState, TSlice
   */

  // 返回最终的、绑定了 API 的 useStore Hook
  return useStore;
}


// -----------------------------------------------------------------------------
// 3. `create` 函数 (主入口)
//    组合 createStoreApi 和 createUseStoreHook
// -----------------------------------------------------------------------------

/**
 * 创建 Zustand 风格的 Store 和对应的 React Hook。
 * @template TState
 * @param {function(
 *   function(TState | function(TState): TState, boolean?): void,
 *   function(): TState,
 *   StoreApi<TState>
 * ): TState} initializer - 用户提供的初始化函数。
 * @returns {UseBoundStore<TState>} 返回绑定了 Store API 的 React Hook。
 *
 * @example
 * const useBearStore = create((set, get) => ({
 *   bears: 0,
 *   increasePopulation: () => set(state => ({ bears: state.bears + 1 })),
 *   removeAllBears: () => set({ bears: 0 }),
 *   getBears: () => get().bears,
 * }));
 */
function create(initializer) {
  // 1. 创建 Store 的核心 API 实例
  const storeApi = createStoreApi(initializer);
  // 2. 基于 Store API 创建对应的 React Hook
  const useStore = createUseStoreHook(storeApi);
  // 3. 返回这个 Hook,它上面附加了 Store API
  return useStore;
}


// -----------------------------------------------------------------------------
// 4. 中间件 (示例)
// -----------------------------------------------------------------------------

/**
 * 简单的日志中间件示例。
 * @param {function} config - Store 的原始初始化函数 (createState)。
 * @returns {function} - 应用了中间件的新的初始化函数。
 */
const loggerMiddleware = (config) => (set, get, api) => {
  // 包装原始的 set 函数
  const loggedSet = (args, replace, actionType) => {
    console.groupCollapsed(`[Logger Middleware] Action: ${actionType || 'Unknown Action'}`);
    console.log('%c Previous State:', 'color: gray', get());
    console.log('%c Action Payload:', 'color: blue', typeof args === 'function' ? '<function>' : args);
    // 调用原始的 set 函数
    set(args, replace, actionType); // 将 actionType 传递下去
    console.log('%c Next State:', 'color: green', get());
    console.groupEnd();
  };

  // 调用原始的 config 函数,但传入我们包装过的 set
  // get 和 api 保持不变
  const initialState = config(loggedSet, get, api);

  console.log('[Logger Middleware] Store initialized with middleware.');

  // 返回初始状态
  return initialState;
};

/**
 * 持久化中间件的极简模拟 (仅 localStorage, 基础功能)。
 * @param {function} config - 原始初始化函数。
 * @param {object} options - 配置项。
 * @param {string} options.name - localStorage 的 key。
 * @param {function(Storage): Storage} [options.storage = () => localStorage] - 存储引擎。
 * @param {function(TState): Partial<TState>} [options.partialize = state => state] - 选择要持久化的部分状态。
 * @returns {function} - 应用了中间件的新的初始化函数。
 */
const persistMiddleware = (config, options) => (set, get, api) => {
    const { name, storage: getStorage = () => localStorage, partialize = state => state } = options;
    let storage;
    try {
        storage = getStorage();
    } catch (e) {
        console.error("[Persist Middleware] Failed to get storage:", e);
        // 如果无法获取存储,则中间件失效,直接调用原始 config
        return config(set, get, api);
    }

    let hasHydrated = false; // 标记是否已从存储恢复

    // 包装 set,使其在更新状态后也更新 localStorage
    const persistedSet = (args, replace, actionType) => {
        // 先调用原始 set 更新内存中的状态
        set(args, replace, actionType);
        // 然后将需要持久化的部分写入存储
        try {
            const stateToPersist = partialize(get());
            storage.setItem(name, JSON.stringify({ state: stateToPersist }));
             console.log(`[Persist Middleware] Saved state to ${name}:`, stateToPersist);
        } catch (e) {
            console.error(`[Persist Middleware] Failed to save state to ${name}:`, e);
        }
    };

    // 调用原始 config 获取初始状态(可能包含 actions)
    const initialState = config(persistedSet, get, api);

    // --- Hydration (从存储恢复状态) ---
    try {
        const storedValue = storage.getItem(name);
        if (storedValue) {
            const persistedState = JSON.parse(storedValue).state;
             console.log(`[Persist Middleware] Rehydrating state from ${name}:`, persistedState);
            // 将恢复的状态合并到初始状态中
            // 注意:这里简单合并,可能覆盖 initialState 中的默认值
            // 真实 Zustand 有更复杂的合并/版本管理逻辑
            Object.assign(initialState, persistedState);
            hasHydrated = true;
             // 用恢复后的完整状态调用一次原始 set,以确保内存状态是最新的
             // 使用 replace=true 确保完全覆盖
             set(initialState, true, 'persist/rehydrate');
        } else {
             console.log(`[Persist Middleware] No state found in ${name} for hydration.`);
        }
    } catch (e) {
        console.error(`[Persist Middleware] Failed to rehydrate state from ${name}:`, e);
    }

    console.log('[Persist Middleware] Store initialized. Hydrated:', hasHydrated);

    // 返回(可能已恢复的)初始状态
    return initialState;
};


// -----------------------------------------------------------------------------
// 5. 使用示例
// -----------------------------------------------------------------------------

console.log("\n--- Example 1: Simple Counter Store ---");
const useCounterStore = create((set, get) => ({
  count: 0,
  increment: () => set(state => ({ count: state.count + 1 }), false, 'INCREMENT'),
  decrement: () => set(state => ({ count: state.count - 1 }), false, 'DECREMENT'),
  getCount: () => get().count, // Action 使用 get
}));

// --- 在组件外部使用 API ---
console.log("Initial count:", useCounterStore.getState().count);
useCounterStore.setState({ count: 10 }, false, 'SET_COUNT_EXTERNAL');
console.log("Count after external set:", useCounterStore.getState().count);

// 订阅
const unsubscribeCounter = useCounterStore.subscribe((newState, oldState) => {
  console.log(`[Subscription Callback] Count changed: ${oldState.count} -> ${newState.count}`);
});

// 触发 action
useCounterStore.getState().increment(); // count: 11
useCounterStore.getState().increment(); // count: 12

// 取消订阅
unsubscribeCounter();
console.log("Unsubscribed from counter store.");
useCounterStore.getState().decrement(); // count: 11 (不会触发上面的回调)
console.log("Final count:", useCounterStore.getState().count);


console.log("\n--- Example 2: Store with Middleware (Logger + Persist) ---");

const useFishStore = create(
  // 应用中间件:从外到内包裹
  loggerMiddleware( // 最外层是 logger
    persistMiddleware( // 内层是 persist
      // 原始的 config 函数
      (set, get) => ({
        fishes: 0,
        bubbles: 10,
        addFish: () => set(state => ({ fishes: state.fishes + 1 }), false, 'ADD_FISH'),
        addBubble: () => set(state => ({ bubbles: state.bubbles + 1 }), false, 'ADD_BUBBLE'),
        eatFish: () => set(state => ({ fishes: Math.max(0, state.fishes -1) }), false, 'EAT_FISH'),
        // 模拟持久化部分状态
        getPersistedData: () => ({ fishes: get().fishes }), // 假设只持久化 fishes
      }),
      // Persist options
      {
        name: 'fish-storage', // localStorage key
        partialize: (state) => ({ fishes: state.fishes }), // 只持久化 fishes
      }
    )
  )
);

// 初始状态(可能会从 localStorage 恢复 fishes)
console.log("Initial fish state:", useFishStore.getState());

// 触发 actions (会经过 logger 和 persist)
useFishStore.getState().addFish();
useFishStore.getState().addBubble(); // bubbles 不会被持久化
useFishStore.getState().addFish();

console.log("Current fish state:", useFishStore.getState());

// 可以在浏览器的 localStorage 查看 'fish-storage' 的值

// 清理 store (移除监听器,但 localStorage 数据还在)
// useFishStore.destroy();


console.log("\n--- Example 3: Simulating React Component Usage ---");

// 模拟一个 React 组件
function CounterComponent() {
  console.log("\n[CounterComponent] Rendering...");

  // 使用 Hook 订阅 count 状态
  // const count = useCounterStore(state => state.count);
  // 使用 Hook 订阅 count 状态 (使用模拟 React)
  const count = useStore(useCounterStore, state => state.count);

  console.log(`[CounterComponent] Selected count: ${count}`);

  // 使用 Hook 获取 action (action 引用是稳定的)
  // const increment = useCounterStore(state => state.increment);
  const increment = useStore(useCounterStore, state => state.increment);

  // 模拟按钮点击
  const simulateClick = () => {
    console.log("[CounterComponent] Simulating button click to increment...");
    increment();
    // 在真实 React 中,状态更新会触发重新渲染
    // 在模拟中,我们需要手动调用组件函数来模拟重渲染
    // CounterComponent(); // 手动模拟重渲染 (如果需要观察后续状态)
  };

  // 模拟组件挂载时的效果
  React.useEffect(() => {
    console.log("[CounterComponent] Mounted. Initial count was:", count);
    // 可以在这里执行一些副作用
    return () => {
      console.log("[CounterComponent] Unmounted.");
    };
  }, []); // 空依赖数组,只在挂载和卸载时运行

  return {
    render: () => `<div>Count: ${count}</div>`,
    simulateClick: simulateClick,
  };
}

// 模拟组件渲染和交互
const myCounter = CounterComponent();
console.log(myCounter.render());
myCounter.simulateClick(); // 触发 increment action
// 观察控制台输出,看状态更新和监听器通知

// 模拟另一个只关心不同状态的组件
function NonCounterComponent() {
    console.log("\n[NonCounterComponent] Rendering...");
    // 这个组件不关心 count,它只订阅一个不存在的状态(或者 store 的其他部分)
    // const other = useCounterStore(state => state.other, shallowCompare); // 假设用了 shallow
    const other = useStore(useCounterStore, state => state.other, shallowCompare);
    console.log(`[NonCounterComponent] Selected other: ${other}`); // 应该是 undefined

     React.useEffect(() => {
        console.log("[NonCounterComponent] Mounted.");
        return () => {
          console.log("[NonCounterComponent] Unmounted.");
        };
      }, []);

    return { render: () => `<div>Other: ${other}</div>` };
}

const myNonCounter = NonCounterComponent();
console.log(myNonCounter.render());

// 再次触发 counter store 的更新
console.log("\nTriggering another increment on counter store...");
useCounterStore.getState().increment();

// 观察输出:
// - CounterComponent 的 useStore Hook 内部应该检测到 slice (count) 变化并通知 React (模拟)。
// - NonCounterComponent 的 useStore Hook 内部应该检测到 slice (other) 没有变化,不会通知 React (模拟)。


// 导出 create 函数,使其可以像库一样使用
// export default create;
// 在这个脚本环境中,我们直接使用 create

代码讲解:

  1. React 模拟 (React 对象):

    • 为了能在非 React 环境(如 Node.js)或没有安装 React 的情况下运行和理解代码,提供了一个极简的 useState, useEffect, useRefuseSyncExternalStore 的模拟。
    • useSyncExternalStore 的 shim 是关键,它展示了如何使用 useStateuseEffect 来订阅外部数据源并在变化时强制更新组件。注意:这个 shim 并不完全等同于 React 官方实现,尤其是在并发模式下的行为保证。
    • shallowCompare: 提供了类似 zustand/shallow 的浅比较函数,用于 useStoreequalityFn
  2. createStoreApi(createState):

    • 这是 Zustand 的核心,完全不依赖 React。

    • 闭包状态 : statelisteners 被保存在函数闭包中。

    • setState(partial, replace, actionType) :

      • 计算 nextStateFragment (处理对象或函数)。
      • 根据 replace 参数决定是浅合并 (Object.assign({}, state, nextStateFragment)) 还是完全替换。注意:浅合并是关键,与 Redux 的 reducer 不同。
      • 异步通知 : 使用 Promise.resolve().then() 将监听器调用推迟到微任务队列,模拟 React 的批处理,并防止在同步代码执行期间触发过多更新。
      • 状态比较 : 只有当 state 引用实际发生变化时(使用 Object.is 比较 previousStatestate),才会通知监听器。
      • 传递 newStatepreviousState 给监听器。
      • actionType 用于调试和中间件。
    • getState() : 直接返回闭包中的 state

    • subscribe(listener) : 将 listener 添加到 listeners Set,并返回一个用于移除监听器的 unsubscribe 函数。

    • destroy() : 清空 listeners Set。

    • api 对象 : 封装核心方法,传递给 createState 并最终附加到 useStore Hook。

    • 初始化 : 调用用户传入的 createState(set, get, api) 来计算初始状态 state

  3. createUseStoreHook(api):

    • 接收 createStoreApi 返回的 api 对象。

    • useStore(selector, equalityFn) Hook:

      • useSyncExternalStore 实现 (首选) :

        • subscribe: 定义如何订阅 store。关键在于内部的 listener,它在 store 更新时运行 selector,并使用 equalityFn 比较新旧 slice,只有在 slice 变化时才调用 React 的 onStoreChange 回调。依赖项包含 selectorequalityFn 很重要。
        • getSnapshot: 定义如何获取当前选定状态的快照。
        • getServerSnapshot: 用于 SSR。
        • 这是连接外部 store 和 React 组件状态的最现代、最正确的方式。
      • (备选) useState/useEffect 实现:

        • 使用 useState 保存组件本地的 slice
        • 使用 useRef 保存 selector, equalityFn, slice 的最新引用,避免 useEffect 闭包问题。
        • useEffect (空依赖数组) 负责在组件挂载时订阅 store,卸载时取消订阅。
        • 内部 listener 在 store 更新时,使用 ref 中的最新 selectorequalityFn 比较新旧 slice,如果变化则调用 setSlice 更新组件状态。
        • 这个版本更容易理解订阅机制,但不如 useSyncExternalStore 精确和高效。
    • API 附加 : 使用 Object.assign(useStore, api)getState, setState 等核心 API 方法附加到 useStore 函数本身,允许在组件外部调用,例如 useCounterStore.getState()

  4. create(initializer):

    • 这是库的主入口函数。
    • 它简单地调用 createStoreApi 创建核心逻辑,然后调用 createUseStoreHook 创建 React Hook,最后返回这个 Hook。
  5. 中间件 (loggerMiddleware, persistMiddleware) :

    • 中间件本质上是高阶函数,它们接收一个 config (即 createState 函数) 并返回一个新的 config 函数。
    • 包装 set : 中间件通常通过包装 set 函数来注入逻辑。例如,loggerMiddleware 在调用原始 set 前后打印日志,persistMiddleware 在调用原始 set 后将状态写入存储。
    • 状态恢复 (Hydration) : persistMiddleware 尝试在初始化时从存储中读取状态,并将其合并到 initialState 中,然后调用一次 set(..., true) 来确保内存状态与恢复的状态一致。
    • 组合: 可以像示例中那样嵌套调用中间件来组合它们的功能。执行顺序通常是从外到内。
  6. 示例:

    • 演示了如何创建简单 store (useCounterStore)。
    • 演示了如何在组件外部使用附加到 Hook 上的 API (getState, setState, subscribe)。
    • 演示了如何组合使用中间件 (useFishStore)。
    • 演示了如何在模拟的 React 组件 (CounterComponent, NonCounterComponent) 中使用 useStore Hook,包括选择器和渲染优化(通过比较 slice)。

与真实 Zustand 的差异和简化:

  • React 版本兼容性 : 这个实现主要依赖 useSyncExternalStore (或其 shim)。真实的 Zustand 对旧版 React 有更完善的兼容处理。
  • 性能优化 : 真实的 Zustand 可能包含更多微优化,例如更精细的批处理、proxy-compare 支持(用于更精确地跟踪对象属性访问以进行比较)。
  • 中间件 API: 真实的 Zustand 中间件 API 可能更复杂或提供更多辅助功能。这里的中间件实现是基础的包装模式。
  • 错误处理: 错误处理相对简单,真实的库可能有更健壮的错误边界和报告机制。
  • 类型系统: 没有包含 TypeScript 类型定义,真实的 Zustand 提供了强大的类型支持。
  • 特性完整性 : 可能缺少一些 Zustand 的高级特性或边缘情况的处理(如瞬态更新 temporal 中间件、复杂的选择器组合等)。
  • proxy-compare: 未实现此高级优化,它允许仅在选择器实际访问的属性发生更改时才重新渲染,即使选择器返回的是一个新对象。
  • Persist 细节 : 模拟的 persist 中间件非常基础,缺少版本迁移、合并策略、异步存储支持等高级功能。

尽管存在简化,这个手写版本准确地反映了 Zustand 的核心设计哲学:利用闭包管理状态,通过订阅模式通知更新,并利用 React Hooks(尤其是 useSyncExternalStore)和选择器/比较函数实现高效、精确的组件渲染。代码量也比较可观,并包含了详细的注释和原理说明。

相关推荐
Java~~16 分钟前
山东大学软件学院项目实训-基于大模型的模拟面试系统-前端美化滚动条问题
前端
掘金安东尼27 分钟前
🧭 前端周刊第411期(2025年4月21日–27日)
前端·javascript·面试
喝养乐多长不高36 分钟前
详细PostMan的安装和基本使用方法
java·服务器·前端·网络协议·测试工具·https·postman
阿珊和她的猫1 小时前
React 与 Vue 的区别:你会选择哪个框架呢
前端·vue.js·react.js
尖椒土豆sss1 小时前
Nuxt3框架入门:第一个简单demo
前端·nuxt.js
浪裡遊1 小时前
前端高频面试题day2
前端·javascript·vue.js
独立开阀者_FwtCoder1 小时前
Vue3 首款 3D 数字孪生编辑器 正式开源!
前端·javascript·vue.js
晓得迷路了1 小时前
栗子前端技术周刊第 78 期 - React Compiler RC、Trae MCP、pnpm 10.9...
前端·javascript·trae
神秘代码行者1 小时前
使用Hash和HTML5的History API实现前端路由
前端·html5