chenglongavatar

chenglong

zustand源码解读-上 - chenyuan

notion image
This article discusses the zustand v4.3.8 library for state management, which is based on the publish-subscribe pattern and can be used outside of React projects. The create function generates a store, which is a closure with exposed APIs for accessing the store. The useSyncExternalStoreWithSelector function binds the store to the view layer, allowing external stores to control page display. The article provides code examples for using zustand in JavaScript and TypeScript projects.

使用方式

zustand是基于发布订阅模式实现的一个状态管理库，可以不局限于仅在react项目中使用，不过对react的支持是官方实现的，使用起来也非常简洁，使用示例如下
// 在js项目中使用，不需要类型
import { create } from "zustand";

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

const useBearStore = create(initStateCreateFunc);

// 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的一模一样
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>;
}

源码主体流程

zustand的核心是将外部store和组件view的交互，交互的核心流程如下图
notion image
先使用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一模一样

简化带注释源码

// 生成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;

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

Copyright © 2024 chenglong

logo