Zustand 源码解析：第一章 external store

查看原文

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()。

发布订阅模式的实现

基础实现

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

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 的状态。更新状态时，可以直接传递下一个状态，也可以传递一个根据上一个状态计算出下一个状态的函数。

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

    state = nextState;

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

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

  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 中，会将新旧状态提供给订阅状态变更的监听者。

  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 中的设计思路如下：

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 和状态放置在一起），像下面这样：

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

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

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 的构造函数：

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。

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 暴露给使用者：

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

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

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

  return useBoundStore;
};

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

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() 做了柯里化。

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