chenyuan

chenyuan

twitter

zustand原碼解讀-上

本文章基於 zustand v4.3.8
微信推送 知乎鏈接
zustand 的4.3.8 tags 鏈接
zustand 的文檔地址

使用方式#

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 基本的使用方式和reduxuseSelector的一模一樣

function BearCounter() {
  const bears = useBearStore((state) => state.bears);
  return <h1>{bears} 在這裡...</h1>;
}

function Controls() {
  const increase = useBearStore((state) => state.increase);
  return <button onClick={increase}>增加一隻</button>;
}

細心的你可能發現 js 和 ts 的使用有不同,ts 是 create<BearState>()(initStateCreateFunc),原因會在下文解釋

源碼主體流程#

zustand的核心是將外部store和組件view的交互,交互的核心流程如下圖
zustand

先使用create函數基於注入的initStateCreateFunc創建一個閉包的store,並暴露對應的subscribesetStategetStatedestory(此 api 將被移除) 這幾個api

借助於react官方提供的useSyncExternalStoreWithSelector可以將storeview層綁定起來,從而實現使用外部的store來控制頁面的展示。

zustand還支持了middleware的能力,採用create(middleware(...args))的形式即可使用對應的middleware

核心代碼詳解#

這部分講解最核心的createuseSyncExternalStoreWithSelector函數

create 函數生成 store#

為了便於閱讀,代碼有刪減

前置知識介紹#

create函數生成的store是一個閉包,通過暴露api的方式實現對store的訪問。

核心代碼在vanilla.tsreact.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一模一樣

簡化帶註釋源碼#

看代碼會發現createStorecreate這兩個函數都是(createState) => createState ? createStoreImpl(createState) : createStoreImpl的形式,翻閱官方文檔的ts guide,會發現官方在 ts 項目裡的調用方式是這樣的:create()(...args)

在文檔代碼示例下方有解釋是為了處理TypeScript/issues/10571而實現的一個 workaround,這也是上文 ts 和 js 使用方式不一致的解答

// 生成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);
    // 取消訂閱
    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,封裝後支持了selectorisEqual

useSyncExternalStore一定需要傳入subscribegetSnapshot兩個函數,返回值是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;
}
載入中......
此文章數據所有權由區塊鏈加密技術和智能合約保障僅歸創作者所有。