Zustand 简介

一个 React 状态管理库，具备如下优点：

小巧：包大小为 1.18 KB；基于 React 的 useSyncExternalStore() hook。
快速：通过 selector，可以减少不必要的重渲染；subscribe 函数允许组件绑定到状态端口，而不会在发生变化时强制重新渲染。当允许直接更改视图时，这会对性能产生巨大影响。
可扩展：灵活的 API；中间件机制。

Zustand 原理简介

Zustand 提供了一种发布订阅模式的实现，即 createStore()。在 React 中，使用 useSyncExternalStore() hook 订阅 store，实现状态更新后，触发组件重新渲染。

出于兼容性考虑以及 selector 实现，Zustand 并未使用 react 库提供的 useSyncExternalStore()，而是使用了 use-sync-external-store 库提供的 useSyncExternalStoreWithSelector()。

发布订阅模式的实现

基础实现

如下是一个简单实现（函数式编程）：

javascript 复制代码

const createStore = () => {
  let state;
  const listeners = new Set();

  const getState = () => state;

  const setState = (newState) => {
    state = newState;

    listeners.forEach((listener) => {
      listener();
    });
  };

  const subscribe = (listener) => {
    listeners.add(listener);

    const unsubscribe = () => listeners.delete(listener);

    return unsubscribe;
  };

  const store = {
    getState,
    setState,
    subscribe,
  };

  return store;
};

const store = createStore();

store 具备状态管理能力，我们可以通过 store 获取状态，修改状态，或者订阅状态变更。

在 createStore 函数内，我们创建了两个变量，一个是 state，一个是 listeners。其中 state 为当前 store 的状态，而 listeners 为订阅状态变更的监听者。

代码中 store 具有三个方法：

getState() ：返回 store 的状态。
setState()：更新 store 的状态，更新状态后，需要通知所有订阅状态变更的监听者。
subscribe()：订阅 store 的状态变更。subscribe() 接受一个 listener 函数作为参数，当状态变更后，会调用 listener() 函数。即我们在 setState() 中更新状态后，遍历当前的所有 listener，并调用。同时需要返回一个 unsubscribe 函数，这个函数的作用是取消订阅。

相较于 Zustand 的实现，上述实现有两个优化点：状态更新和初始状态

优化点一：状态更新

在 Zustand 中，store 的 setState() 方法可以更新 store 的状态。更新状态时，可以直接传递下一个状态，也可以传递一个根据上一个状态计算出下一个状态的函数。

javascript 复制代码

  const setState = (newState) => {
    // 更新状态时，可以直接传递下一个状态，也可以传递一个根据上一个状态计算出下一个状态的函数。
    let nextState = typeof newState === "function" ? newState(state) : newState;

    state = nextState;

    listeners.forEach((listener) => {
      listener();
    });
  };

除此之外，Zustand 中对于是引用类型的状态值，会进行浅层合并。首先会通过 Object.is() 比较下一状态与当前状态，如果相等，忽略更新。同时支持第二个参数 replace，用于指定是否直接替换状态，而不是进行浅层合并。源码中使用 Object.assign() 实现浅层合并。

javascript 复制代码

  const setState = (newState, replace) => {
    let nextState = typeof newState === "function" ? newState(state) : newState;

    // 首先会通过 `Object.is()` 比较下一状态与当前状态，如果相等，忽略更新。
    if (!Object.is(nextState, state)) {
      state =
        // replace 不为 `undefined` 或下一状态值为原始类型的值
        replace ?? (typeof nextState !== "object" || typeof nextState === null)
          ? nextState
          : Object.assign({}, state, nextState);

      listeners.forEach((listener) => {
        listener();
      });
    }
  };

同时，Zustand 中，会将新旧状态提供给订阅状态变更的监听者。

javascript 复制代码

  const setState = (newState, replace) => {
    let nextState = typeof newState === "function" ? newState(state) : newState;

    if (!Object.is(nextState, state)) {
      const previousState = state;

      state =
        replace ?? (typeof nextState !== "object" || typeof nextState === null)
          ? nextState
          : Object.assign({}, state, nextState);

      listeners.forEach((listener) => {
        // 将新旧状态提供给订阅状态变更的监听者
        listener(state, previousState);
      });
    }
  };

优化点二：初始状态

对于我们目前的实现，创建的 store 的状态均是 undefined。我们需要提供一种方式，让使用者可以定义 store 的初始状态。这里我们可以参考 React 的 useState hook 的设计，useState() 的第一个参数是 initialState，即初始状态。initialState 可以是任何类型的值，但如果是函数时，initialState 会被视为初始化函数，其函数调用的结果为初始值。这种设计提供了很好的可扩展性，以及解决了 React 中特有的问题：避免重新创建初始状态。

Zustand 中的设计思路如下：

javascript 复制代码

const createStore = (createInitialState) => {
  let state;
  const listeners = new Set();
  
  // ...

  const store = {
    getState,
    setState,
    subscribe,
  };

  // 调用初始化函数，初始化状态
  state = createInitialState();

  return store;
};

const store = createStore(() => {
  return {
    count: 1,
  };
});

在 Zustand 中支持 action 直接添加到 store 中（让 action 和状态放置在一起），像下面这样：

javascript 复制代码

const useBoundStore = create((set) => ({
  count: 1,
  inc: () => set((state) => ({ count: state.count + 1 })),
}))

为了实现这个功能特性，我们需要将 store 相关的 API 作为初始化函数的参数，提供给使用者。

javascript 复制代码

const createStore = (initialState) => {
  let state;
  const listeners = new Set();

  // ...

  const store = {
    getState,
    setState,
    subscribe,
  };

  // 将 store 相关的 API 作为初始化函数的参数，提供给使用者
  state = initialState(setState, getState, store);

  return store;
};

const store = createStore((set) => {
  return {
    count: 1,
    resetCount: () => set(1),
  };
});

useBoundStore 的构造函数 create 的实现

Zustand 提供的 create() 函数，是一个自定义 hook 的构造函数：

javascript 复制代码

const create = (createInitialState) => {
  // 创建 store
  const store = createStore(createInitialState);

  const useBoundStore = () => {
    // ...
  };

  return useBoundStore;
};

该函数会创建 store，并返回一个自定义 hook，这个自定义 hook 是对 useSyncExternalStoreWithSelector() 的封装。其来自 use-sync-external-store 库，相比 React 中的 useSyncExternalStore() 有更好的兼容性和特性支持。

特性支持：

selector: 从 store 中选择状态（计算新状态），优化状态更新粒度，实现渲染优化。
isEqual: 自定义相等性判断。

返回的自定义 hook 对 useSyncExternalStoreWithSelector() hook 封装了一层，让使用者不再自己手动创建 store。

javascript 复制代码

const create = (createInitialState) => {
  const store = createStore(createInitialState);

  const useBoundStore = (selector = () => store, equalityFn) => {
    // 调用 useSyncExternalStoreWithSelector()
    return useSyncExternalStoreWithSelector(
      store.subscribe,
      store.getState,
      store.getState,
      selector,
      equalityFn
    );
  };

  return useBoundStore;
};

同时 Zustand 有将 store 相关的 API 暴露给使用者：

javascript 复制代码

const create = (createInitialState) => {
  const store = createStore(createInitialState);

  const useBoundStore = () => {
    // ...
  };

  // 将 store 相关的 API 暴露给使用者
  Object.assign(useBoundStore, store);

  return useBoundStore;
};

Zustand 源码中将 useSyncExternalStoreWithSelector() 单独封装为 useStoreWithEqualityFn()，以优化代码复用和代码扩展性。

javascript 复制代码

const useStoreWithEqualityFn = (store, selector = api.getState, equalityFn) => {
  return useSyncExternalStoreWithSelector(
    store.subscribe,
    store.getState,
    store.getState,
    selector,
    equalityFn
  );
};

const create = (createInitialState) => {
  const store = createStore(createInitialState);

  const useBoundStore = (selector = store.getState, equalityFn) => {
    return useStoreWithEqualityFn(store, selector, equalityFn);
  };

  Object.assign(useBoundStore, store);

  return useBoundStore;
};

为了在调用 create() （具有多个泛型的函数）时，允许跳过某些泛型。即想要注释状态（第一个类型参数），但允许推断其他参数时。Zustand 对 create() 做了柯里化。

javascript 复制代码

const createImpl = (createInitialState) => {
  const store = createStore(createInitialState);

  const useBoundStore = (selector = store.getState, equalityFn) => {
    return useStoreWithEqualityFn(store, selector, equalityFn);
  };

  Object.assign(useBoundStore, store);

  return useBoundStore;
};

// 如果没有传递初始化函数，直接返回
const create = (createInitialState) => {
  return createInitialState ? createImpl(createInitialState) : createImpl;
};

Zustand 源码解析：第一章 external store

Zustand 简介

Zustand 原理简介

发布订阅模式的实现

基础实现

优化点一：状态更新

优化点二：初始状态

useBoundStore 的构造函数 create 的实现