zustand 源码解析

信息

此文章基于 zustand v4.3.8

zustand 的4.3.8 tags 链接， zustand 的文档地址

使用方式

zustand是基于发布订阅模式实现的一个状态管理库，可以不局限于仅在react项目中使用，不过对react的支持是官方实现的，使用起来也非常简洁，使用示例如下

js

// 在js项目中使用，不需要类型
import { create } from "zustand";

const initStateCreateFunc = (set) => ({
  bears: 0,
  increase: (by) => set((state) => ({ bears: state.bears + by })),
});

const useBearStore = create(initStateCreateFunc);

ts

// ts项目，需要类型提示
import { create } from "zustand";

interface BearState {
  bears: number;
  increase: (by: number) => void;
}

const initStateCreateFunc = (set) => ({
  bears: 0,
  increase: (by) => set((state) => ({ bears: state.bears + by })),
});

const useBearStore = create<BearState>()(initStateCreateFunc);

如上文代码，在调用create函数后，会生成一个useStore的 hook，这个 hook 基本的使用方式和redux的useSelector的一模一样

在组件中引入useStore

function BearCounter() {
  const bears = useBearStore((state) => state.bears);
  return <h1>{bears} around here...</h1>;
}

function Controls() {
  const increase = useBearStore((state) => state.increase);
  return <button onClick={increase}>one up</button>;
}

提示

细心地你可能发现 js 和 ts 的使用有不同，ts 是 create<BearState>()(initStateCreateFunc)，原因会在下文解释

源码主体流程

zustand的核心是将外部store和组件view的交互，交互的核心流程如下图

先使用create函数基于注入的initStateCreateFunc创建一个闭包的store，并暴露对应的subscribe、setState、getState、destory(此 api 将被移除)这几个api

借助于react官方提供的useSyncExternalStoreWithSelector可以将store和view层绑定起来，从而实现使用外部的store来控制页面的展示。

zustand还支持了middleware的能力，采用create(middleware(...args))的形式即可使用对应的middleware

核心代码详解

这部分讲解最核心的create和useSyncExternalStoreWithSelector函数

create 函数生成 store

信息

为了便于阅读，代码有删减

前置知识介绍

create函数生成的store是一个闭包，通过暴露api的方式实现对store的访问。

核心代码在vanilla.ts和react.ts这两个文件中，vanilla.ts里实现了一个完整的有pub-sub能力的store, 不需要依赖于react即可使用。

react.ts里基于useSyncExternalStoreWithSelector实现了一个useStore的 hook，在组件里调用create返回的函数时会将store和组件绑定起来，而这个绑定就是useStore实现的 这个useSyncExternalStoreWithSelector会在下一小节讲述。

create 运行流程

在create函数调用的时候，先使用vanilla.ts导出的createStore生成store，然后定义一个useBoundStore函数，返回值是useStore(api, selector, equalityFn)，然后把createStore返回的api注入useBoundStore上，然后返回useBoundStore. 这个useBoundStore的使用方式和useSelector一模一样

简化带注释源码

提示

看代码会发现createStore和create这两个函数都是(createState) => createState ? createStoreImpl(createState) : createStoreImpl的形式，翻阅官方文档的ts guide，会发现官方在 ts 项目里的调用方式是这样的：create()(...args)

在文档代码示例下方有解释是为了处理TypeScript/issues/10571而实现的一个 walkaround，这也是上文 ts 和 js 使用方式不一致的解答

vanilla.ts

vanilla.ts

// 生成store闭包，并返回api
// createState是使用者在创建store时传入的一个函数
const createStoreImpl = (createState) => {
  type TState = ReturnType<typeof createState>;
  type Listener = (state: TState, prevState: TState) => void;
  // 这里的state就是store，是个闭包，通过暴露的api访问
  let state: TState;
  const listeners: Set<Listener> = new Set();

  // setState的partial参数支持对象和函数，replace指明是全量替换store还是merge
  // 更新是浅比较
  const setState = (partial, replace) => {
    const nextState = typeof partial === "function" ? partial(state) : partial;
    // 只有在相等的时候才更新，然后触发listener
    if (!Object.is(nextState, state)) {
      const previousState = state;
      state =
        replace ?? typeof nextState !== "object"
          ? (nextState as TState)
          : Object.assign({}, state, nextState);
      listeners.forEach((listener) => listener(state, previousState));
    }
  };

  const getState = () => state;

  const subscribe = (listener) => {
    listeners.add(listener);
    // Unsubscribe
    return () => listeners.delete(listener);
  };

  // destory之后将被去掉，不用看
  const destroy: StoreApi<TState>["destroy"] = () => {
    if (import.meta.env?.MODE !== "production") {
      console.warn(
        "[DEPRECATED] The `destroy` method will be unsupported in a future version. Instead use unsubscribe function returned by subscribe. Everything will be garbage-collected if store is garbage-collected."
      );
    }
    listeners.clear();
  };

  const api = { setState, getState, subscribe, destroy };
  // 这里就是官方示例里的set,get,api
  state = createState(setState, getState, api);
  return api as any;
};

// 调用createStore的时候理论上createState函数是一定存在的
// 但是为了ts类型定义，createStore<T>()(()=>{}) 所以会出现手动调用空值的情况
export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore;

react.ts

react.ts

export function useStore<TState, StateSlice>(
  api: WithReact<StoreApi<TState>>,
  selector: (state: TState) => StateSlice = api.getState as any,
  equalityFn?: (a: StateSlice, b: StateSlice) => boolean
) {
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getState,
    selector,
    equalityFn
  );
  useDebugValue(slice);
  return slice;
}

const createImpl = (createState) => {
  if (
    import.meta.env?.MODE !== "production" &&
    typeof createState !== "function"
  ) {
    console.warn(
      "[DEPRECATED] Passing a vanilla store will be unsupported in a future version. Instead use `import { useStore } from 'zustand'`."
    );
  }
  // 直接注入自定义的store不会注入api，需要自己在注入的store里自行实现
  const api =
    typeof createState === "function" ? createStore(createState) : createState;

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

  Object.assign(useBoundStore, api);

  return useBoundStore;
};

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

useSyncExternalStoreWithSelector 解析

zustand的核心代码如此简洁，一大原因就是使用了useSyncExternalStoreWithSelector，这个是react官方出的use-sync-external-store/shim/with-selector包，之所以出这个包，是因为react在提出useSyncExternalStore这个 hook 后，在react v18版本做了重新实现，有破坏性更新。为了兼容性考虑出了这个包。

话不多说，上源码

这个实现其实是基于官方的useSyncExternalStore做的一个封装，官方 hook 不支持传入selector，封装后支持了selector和isEqual。

useSyncExternalStore一定需要传入subscribe和getSnapshot两个函数，返回值是getSnapshot的返回结果。react会给subscribe注入一个callback函数，当外部store变化的时候，一定要手动的调用callback，通知react外部store变化了，需要它重新调用getSnapshot获取最新的状态，如果状态改变了就触发re-render，否则不re-render

useSyncExternalStoreWithSelector的优化主要是允许从一个大store中取出组件所用到的部分，同时借助isEqual来减少re-render的次数

export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
  subscribe: (() => void) => () => void,
  getSnapshot: () => Snapshot,
  getServerSnapshot: void | null | (() => Snapshot),
  selector: (snapshot: Snapshot) => Selection,
  isEqual?: (a: Selection, b: Selection) => boolean,
): Selection {
  // Use this to track the rendered snapshot.
  const instRef = useRef<
    | {
        hasValue: true,
        value: Selection,
      }
    | {
        hasValue: false,
        value: null,
      }
    | null,
  >(null);
  let inst;
  if (instRef.current === null) {
    inst = {
      hasValue: false,
      value: null,
    };
    instRef.current = inst;
  } else {
    inst = instRef.current;
  }

  /**
   * zustand使用的时候采用的是useStore(selector)的形式，每次re-render都会获得一个新的selector
   * 所以getSelection在re-render后都是新的，但是因为有instRef.current以及isEqual
   * 当isEqual的时候返回instRef.current缓存的值，也就是getSelection的返回值不变
   * 不会再次re-render，减少了re-render的次数
   *  */
  const [getSelection, getServerSelection] = useMemo(() => {
    // Track the memoized state using closure variables that are local to this
    // memoized instance of a getSnapshot function. Intentionally not using a
    // useRef hook, because that state would be shared across all concurrent
    // copies of the hook/component.
    let hasMemo = false;
    let memoizedSnapshot;
    let memoizedSelection: Selection;
    const memoizedSelector = (nextSnapshot: Snapshot) => {
      if (!hasMemo) {
        // The first time the hook is called, there is no memoized result.
        hasMemo = true;
        memoizedSnapshot = nextSnapshot;
        const nextSelection = selector(nextSnapshot);
        if (isEqual !== undefined) {
          // Even if the selector has changed, the currently rendered selection
          // may be equal to the new selection. We should attempt to reuse the
          // current value if possible, to preserve downstream memoizations.
          if (inst.hasValue) {
            const currentSelection = inst.value;
            if (isEqual(currentSelection, nextSelection)) {
              memoizedSelection = currentSelection;
              return currentSelection;
            }
          }
        }
        memoizedSelection = nextSelection;
        return nextSelection;
      }

      // We may be able to reuse the previous invocation's result.
      const prevSnapshot: Snapshot = (memoizedSnapshot: any);
      const prevSelection: Selection = (memoizedSelection: any);

      if (is(prevSnapshot, nextSnapshot)) {
        // The snapshot is the same as last time. Reuse the previous selection.
        return prevSelection;
      }

      // The snapshot has changed, so we need to compute a new selection.
      const nextSelection = selector(nextSnapshot);

      // If a custom isEqual function is provided, use that to check if the data
      // has changed. If it hasn't, return the previous selection. That signals
      // to React that the selections are conceptually equal, and we can bail
      // out of rendering.
      if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
        return prevSelection;
      }

      memoizedSnapshot = nextSnapshot;
      memoizedSelection = nextSelection;
      return nextSelection;
    };
    // Assigning this to a constant so that Flow knows it can't change.
    const maybeGetServerSnapshot =
      getServerSnapshot === undefined ? null : getServerSnapshot;
    const getSnapshotWithSelector = () => memoizedSelector(getSnapshot());
    const getServerSnapshotWithSelector =
      maybeGetServerSnapshot === null
        ? undefined
        : () => memoizedSelector(maybeGetServerSnapshot());
    return [getSnapshotWithSelector, getServerSnapshotWithSelector];
  }, [getSnapshot, getServerSnapshot, selector, isEqual]);

  const value = useSyncExternalStore(
    subscribe,
    getSelection,
    getServerSelection,
  );

  useEffect(() => {
    inst.hasValue = true;
    inst.value = value;
  }, [value]);

  useDebugValue(value);
  return value;
}

zustand 插件介绍

本文是对前文的一个补充，以persist插件为例介绍zustand的中间件系统

zustand的核心代码实现的非常简洁，功能特性也相对较少。如果想要有更多的特性就需要自行实现或者通过引入中间件的方式使用现有的轮子。

zustand 中间件的本质

如官方文档middlware介绍，zustand的中间件实际上是一个高阶函数，它的入参和create函数相同，都是createInitState类的函数，但是不同的是它的返回值仍然是一个createInitState类的函数，本质上是对createInitState做了一层包裹，注入特定的逻辑，实现对createInitState的改写。

import { create } from "zustand";

const createInitState = (set) => ({
  bees: false,
  setBees: (input) => set((state) => void (state.bees = input)),
});
// no middleware
const useStore = create(createInitState);

// with middleware
const useStoreWithMiddleware = create(
  middleware1(middleware2(createInitState))
);

persist 中间件源码详解

persist 中间件的详细介绍可以看persist

入参是createInitState和options，其中options.name必传，其他的都是可选项，options的默认值是这些

defaultOptions = {
    storage: createJSONStorage<S>(() => localStorage),
    // 决定保存state里的哪些数据
    partialize: (state: S) => state,
    version: 0,
    // 决定如何merge现有数据和保存在storage里的的数据
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
}

可以看到最开始的时候调用了createJSONStorage函数来生成 storage, 从我的理解，这个函数的作用是做一个胶水层，使得getItem支持async storage，核心在于

const str = (storage as StateStorage).getItem(name) ?? null;
// support async storage
if (str instanceof Promise) {
  return str.then(parse);
}
return parse(str);

从createJSONStorage函数可以看出 persist 中间件最开始应该是基于localStorage进行的封装，后来进行的拓展，数据的存储都需要使用 string。可以看到存在这样一个问题： setItem并未支持async storage，没有做await的操作，调用后不管是否完成，这与getItem不一致。不过换个思路，其实也没必要保证setItem完成，因为一般不会在setItem后立刻执行getItem，await的话会损失一部分性能。但是从我的角度还是await下比较好。

createJSONStorage源码见下：

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

export function createJSONStorage<S>(
  getStorage: () => StateStorage,
  options?: JsonStorageOptions
): PersistStorage<S> | undefined {
  let storage: StateStorage | undefined;
  try {
    storage = getStorage();
  } catch (e) {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
    return;
  }
  const persistStorage: PersistStorage<S> = {
    getItem: (name) => {
      const parse = (str: string | null) => {
        if (str === null) {
          return null;
        }
        return JSON.parse(str, options?.reviver) as StorageValue<S>;
      };
      const str = (storage as StateStorage).getItem(name) ?? null;
      // 支持 async storage
      if (str instanceof Promise) {
        return str.then(parse);
      }
      return parse(str);
    },
    setItem: (name, newValue) =>
      (storage as StateStorage).setItem(
        name,
        JSON.stringify(newValue, options?.replacer)
      ),
    removeItem: (name) => (storage as StateStorage).removeItem(name),
  };
  return persistStorage;
}

下面的是核心源码(hydrate函数的内容被省略了，在下一步介绍。 )，可以看到在 zustand 原有的 api 上挂在了一个api.persist用来暴露 persist 的 api，最终的返回值也是createInitState(也就是 config)的返回值 initState

可以看到options.skipHydration和hydrate函数，这个函数的作用是将保存的 state 和现有的 state merge 到一起，就类似于SSR的水和，将事件挂载到相应的 dom 上。

const newImpl = (config, baseOptions) => (set, get, api) => {
  // 这里的S其实就是上文的createInitState的返回值的类型
  type S = ReturnType<typeof config>;
  let options = {
    storage: createJSONStorage<S>(() => localStorage),
    partialize: (state: S) => state,
    version: 0,
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
    // above are default configs
    ...baseOptions,
  };

  let hasHydrated = false;
  const hydrationListeners = new Set<PersistListener<S>>();
  const finishHydrationListeners = new Set<PersistListener<S>>();
  let storage = options.storage;

  // 没有storage，就不保存，直接返回config(..args)也就是initState
  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
        );
        set(...args);
      },
      get,
      api
    );
  }

  // set partialized item
  const setItem = (): void | Promise<void> => {
    // 只保存partialize处理过的state
    const state = options.partialize({ ...get() });
    return (storage as PersistStorage<S>).setItem(options.name, {
      state,
      version: options.version,
    });
  };

  const savedSetState = api.setState;

  // 替换新的setState，每次更新后都种入storage
  api.setState = (state, replace) => {
    savedSetState(state, replace);
    void setItem();
  };

  const configResult = config(
    /**
     *
    这一步是为了让createState函数里注入的setState和api.setState等价
    因为是这样注入的，更新api.setState不会影响注入的set函数
    const api = { setState, getState, subscribe, destroy }
    state = createState(setState, getState, api)
     */
    (...args) => {
      // 这里的set === savedSetState
      set(...args);
      void setItem();
    },
    get,
    api
  );

  // a workaround to solve the issue of not storing rehydrated state in sync storage
  // the set(state) value would be later overridden with initial state by create()
  // to avoid this, we merge the state from localStorage into the initial state.
  let stateFromStorage: S | undefined;

  // rehydrate initial state with existing stored state
  const hydrate = () => {
    ...
  };

  (api as StoreApi<S> & StorePersist<S, S>).persist = {
    setOptions: (newOptions) => {
      options = {
        ...options,
        ...newOptions,
      };

      if (newOptions.storage) {
        storage = newOptions.storage;
      }
    },
    clearStorage: () => {
      storage?.removeItem(options.name);
    },
    getOptions: () => options,
    rehydrate: () => hydrate() as Promise<void>,
    hasHydrated: () => hasHydrated,
    onHydrate: (cb) => {
      hydrationListeners.add(cb);

      return () => {
        hydrationListeners.delete(cb);
      };
    },
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb);

      return () => {
        finishHydrationListeners.delete(cb);
      };
    },
  };

  if (!options.skipHydration) {
    hydrate();
  }

  return stateFromStorage || configResult;
};

这里介绍hydrate函数部分，主要流程就是从 storage 中读取到初始的状态，通过toThenable函数将非async storage的值也转换成promisify的类型，统一函数的调用形式。先判断获取的值是否需要migrate，然后进行merge，再进行暴露的与hydrate有关的函数的调用。

const hydrate = () => {
  if (!storage) return;

  // On the first invocation of 'hydrate', state will not yet be defined (this is
  // true for both the 'asynchronous' and 'synchronous' case). Pass 'configResult'
  // as a backup  to 'get()' so listeners and 'onRehydrateStorage' are called with
  // the latest available state.

  hasHydrated = false;
  // 在没有skipHydration的时候，调用hydrate的时候initState还没有生成，get()结果是undefined，所以需要使用前置生成的configResult(还没有和保存的值merge)
  hydrationListeners.forEach((cb) => cb(get() ?? configResult));

  const postRehydrationCallback =
    options.onRehydrateStorage?.(get() ?? configResult) || undefined;

  // bind is used to avoid `TypeError: Illegal invocation` error
  return toThenable(storage.getItem.bind(storage))(options.name)
    .then((deserializedStorageValue) => {
      // 此步为了实现根据version进行旧数据的迁移
      if (deserializedStorageValue) {
        if (
          typeof deserializedStorageValue.version === "number" &&
          deserializedStorageValue.version !== options.version
        ) {
          if (options.migrate) {
            return options.migrate(
              deserializedStorageValue.state,
              deserializedStorageValue.version
            );
          }
          console.error(
            `State loaded from storage couldn't be migrated since no migrate function was provided`
          );
        } else {
          return deserializedStorageValue.state;
        }
      }
    })
    .then((migratedState) => {
      // 这步进行merge
      stateFromStorage = options.merge(
        migratedState as S,
        get() ?? configResult
      );

      set(stateFromStorage as S, true);
      return setItem();
    })
    .then(() => {
      // TODO: In the asynchronous case, it's possible that the state has changed
      // since it was set in the prior callback. As such, it would be better to
      // pass 'get()' to the 'postRehydrationCallback' to ensure the most up-to-date
      // state is used. However, this could be a breaking change, so this isn't being
      // done now.
      postRehydrationCallback?.(stateFromStorage, undefined);

      // It's possible that 'postRehydrationCallback' updated the state. To ensure
      // that isn't overwritten when returning 'stateFromStorage' below
      // (synchronous-case only), update 'stateFromStorage' to point to the latest
      // state. In the asynchronous case, 'stateFromStorage' isn't used after this
      // callback, so there's no harm in updating it to match the latest state.
      stateFromStorage = get();
      hasHydrated = true;
      finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S));
    })
    .catch((e: Error) => {
      postRehydrationCallback?.(undefined, e);
    });
};

toThenable函数源码

const toThenable =
  <Result, Input>(
    fn: (input: Input) => Result | Promise<Result> | Thenable<Result>
  ) =>
  (input: Input): Thenable<Result> => {
    try {
      const result = fn(input);
      if (result instanceof Promise) {
        return result as Thenable<Result>;
      }
      return {
        then(onFulfilled) {
          return toThenable(onFulfilled)(result as Result);
        },
        catch(_onRejected) {
          return this as Thenable<any>;
        },
      };
    } catch (e: any) {
      return {
        then(_onFulfilled) {
          return this as Thenable<any>;
        },
        catch(onRejected) {
          return toThenable(onRejected)(e);
        },
      };
    }
  };