🔍

useStateはどう作られているのか、Reactのソースを読解してみた

2022/03/06に公開

概要

所属しているコミュニティで、「普段読まないOSSのコードを読んでみよう」という企画があり、それに参加した際にReactのuseStateに関するコードを読みました。

ちょっと勉強になった気がするので、備忘録がてらメモに残します。

https://github.com/facebook/react

useState

useStateはReactHooks.jsというところで定義されていました。

ReactHooks.js
export function useState<S>(
  initialState: (() => S) | S,
): [S, Dispatch<BasicStateAction<S>>] {
  const dispatcher = resolveDispatcher();
  return dispatcher.useState(initialState);
}

うーん。短いですが、なんのこっちゃですね。一つ一つみていきましょう。

引数と返り値の型

export function useState<S>(
  initialState: (() => S) | S,
): [S, Dispatch<BasicStateAction<S>>]

引数はinitialStateで、その型は () => S もしくは S。SはGeneric Typeですね。
そしてuseState自体の型はarrayであり、1つ目の要素がS、2つ目の要素が Dispatch<BasicStateAction<S>> です。

Dispatch<BasicStateAction<S>> って何?ということで、DispatchとBasicStateActionの定義元を確認しましょう。

type BasicStateAction<S> = (S => S) | S;
type Dispatch<A> = A => void;

ふむ。それぞれこういう型のようです。

つまり Dispatch<BasicStateAction<S>> というのは、 ((S => S) | S) => void ですね。
1つ目の要素と同じ型の値を引数として、返り値がvoidとなる関数が2つ目の要素となるようです。
馴染み深い、setStateの型ですね。

ちなみに、私はあまり使ったことがないのですが、setStateの引数については、1つ目の要素と同じ型の値を、引数かつ返り値とした関数とすることも可能なようです。
日本語で書くとややこしいですが、つまりこういう書き方です。

setState(count => count + 1)

検索するとたくさんヒットしたので、私が経験不足のだけだと思いますが、個人的には関数も引数にできるんだな、ということで少し驚きでした。

中身

さて、ようやくuseStateの中身です。
ここでもう一度定義をみてみましょう。

export function useState<S>(
  initialState: (() => S) | S,
): [S, Dispatch<BasicStateAction<S>>] {
  const dispatcher = resolveDispatcher();
  return dispatcher.useState(initialState);
}

といっても中身は2行しかありませんね。
resolveDispatcherというのが主に仕事をしているようです。

しかし、resolveDispatcherからuseStateというメソッドを引き出して、それを返り値としていますが、これではさっぱり意味がわかりません。

なので、resolveDispatcherの定義元を見にいきましょう。

resolveDispatcher

同じくReactHook.jsに定義されていました。

ReactHook.js
function resolveDispatcher() {
  const dispatcher = ReactCurrentDispatcher.current;
  if (__DEV__) {
    if (dispatcher === null) {
      console.error(
        'Invalid hook call. Hooks can only be called inside of the body of a function component. This could happen for' +
          ' one of the following reasons:\n' +
          '1. You might have mismatching versions of React and the renderer (such as React DOM)\n' +
          '2. You might be breaking the Rules of Hooks\n' +
          '3. You might have more than one copy of React in the same app\n' +
          'See https://reactjs.org/link/invalid-hook-call for tips about how to debug and fix this problem.',
      );
    }
  }
  // Will result in a null access error if accessed outside render phase. We
  // intentionally don't throw our own error because this is in a hot path.
  // Also helps ensure this is inlined.
  return ((dispatcher: any): Dispatcher);
}

うーむ、__DEV__というのが開発環境かどうかを見るグローバル変数っぽいので、本番に関係するコードは実質2行くらいですね。

そしてdispatcher、もといReactCurrentDispatcher.currentが引数である関数がreturnされています。

全然わからないので、ReactCurrentDispatcherをさらに調べていきましょう。

ReactCurrentDispatcher

ReactCurrentDispatcher.js
const ReactCurrentDispatcher = {
  /**
   * @internal
   * @type {ReactComponent}
   */
  current: (null: null | Dispatcher),
};

ほとんど何も書かれていませんでした。スカされ続けています。
どうやらReactCurrentDispatcher.currentに実質的な値を格納している箇所が他でありそうです。

RenderWithHooks

こちらの記事を参考に、その場所を(たぶん)特定しました。ReactFiberHooks.new.jsのRenderWithHooksという関数です。
https://sbfl.net/blog/2019/02/09/react-hooks-usestate/

というより、こちらの記事が本記事と趣旨がほとんど同じで、かつuseStateの内部についてとても分かりやすく説明をしてくださっているので、もはやここからはこちらの記事を読んだ方がよいかもしれません。

と言いつつも、ここでばたっと終えても不親切なので、ここからは上記記事を参考にさせていただきながら進めることにします。

話は戻り、該当箇所はここです。

ReactFiberHooks.new.js
if (__DEV__) {
    if (current !== null && current.memoizedState !== null) {
      ReactCurrentDispatcher.current = HooksDispatcherOnUpdateInDEV;
    } else if (hookTypesDev !== null) {
      // This dispatcher handles an edge case where a component is updating,
      // but no stateful hooks have been used.
      // We want to match the production code behavior (which will use HooksDispatcherOnMount),
      // but with the extra DEV validation to ensure hooks ordering hasn't changed.
      // This dispatcher does that.
      ReactCurrentDispatcher.current = HooksDispatcherOnMountWithHookTypesInDEV;
    } else {
      ReactCurrentDispatcher.current = HooksDispatcherOnMountInDEV;
    }
  } else {
    ReactCurrentDispatcher.current =
      current === null || current.memoizedState === null
        ? HooksDispatcherOnMount
        : HooksDispatcherOnUpdate;
  }

if(__DEV__)内は関係ないのでスルーですね。else内を見ていきましょう。

currentかcurrent.memoizedStateが値を持っているかによって、HooksDispatcherOnMountかHooksDispatcherOnUpdateのいずれかをReactCurrentDispatcher.currentに格納しています。

currentとは?ということですが、ざっくり、DOMに既に反映している情報を持っている内部インスタンスのようです。

と言いつつよくわかっていませんが、とりあえずcurrentに関する条件で「初回レンダリングかどうか」を見ているようですね。

今回は初回に焦点を当てて読み進めていきましょう。ということでHooksDispatcherOnMountに移ります。

HooksDispatcherOnMount

ReactFiberHooks.new.js
const HooksDispatcherOnMount: Dispatcher = {
  readContext,

  useCallback: mountCallback,
  useContext: readContext,
  useEffect: mountEffect,
  useImperativeHandle: mountImperativeHandle,
  useLayoutEffect: mountLayoutEffect,
  useInsertionEffect: mountInsertionEffect,
  useMemo: mountMemo,
  useReducer: mountReducer,
  useRef: mountRef,
  useState: mountState,
  useDebugValue: mountDebugValue,
  useDeferredValue: mountDeferredValue,
  useTransition: mountTransition,
  useMutableSource: mountMutableSource,
  useSyncExternalStore: mountSyncExternalStore,
  useId: mountId,

  unstable_isNewReconciler: enableNewReconciler,
};

こうなっています。
ここで、useStateにはmountStateという関数が割り当てられていることがわかります。
つまり、初回レンダリングに関しては、mountStateがuseStateの実質であると言えそうです。

さて、探索もそろそろ終盤です。

mountState

ReactFiberHooks.new.js
function mountState<S>(
  initialState: (() => S) | S,
): [S, Dispatch<BasicStateAction<S>>] {
  const hook = mountWorkInProgressHook();
  if (typeof initialState === 'function') {
    // $FlowFixMe: Flow doesn't like mixed types
    initialState = initialState();
  }
  hook.memoizedState = hook.baseState = initialState;
  const queue: UpdateQueue<S, BasicStateAction<S>> = {
    pending: null,
    interleaved: null,
    lanes: NoLanes,
    dispatch: null,
    lastRenderedReducer: basicStateReducer,
    lastRenderedState: (initialState: any),
  };
  hook.queue = queue;
  const dispatch: Dispatch<
    BasicStateAction<S>,
  > = (queue.dispatch = (dispatchSetState.bind(
    null,
    currentlyRenderingFiber,
    queue,
  ): any));
  return [hook.memoizedState, dispatch];
}

returnの値から見ていきましょう。
[hook.memoizedState, dispatch]という配列を返しています。
いつも使っている、[state, setState]の形になんだか似ていますね。

この直感を証明すべく、引き続き読み進めていきます。

hook.memoizedStateには、initialStateが渡されているようですね。

一方で、dispatchという関数はなんでしょう。
一つ上の行で、dispatchSetState関数にcurrentlyRenderingFiber, queueなどをbindしたものであることがわかります。

もう少しdispatchSetStateを詳しく見ていきます。

dispatchSetState(の概要)

ReactFiberHooks.new.js
function dispatchSetState<S, A>(
  fiber: Fiber,
  queue: UpdateQueue<S, A>,
  action: A,
) {
  (省略)
  const lane = requestUpdateLane(fiber);

  const update: Update<S, A> = {
    lane,
    action,
    hasEagerState: false,
    eagerState: null,
    next: (null: any),
  };

  if (isRenderPhaseUpdate(fiber)) {
    enqueueRenderPhaseUpdate(queue, update);
  } else {
    enqueueUpdate(fiber, queue, update, lane);
    (省略)
   }

結構長かったので、開発環境上での動きについてはカットしました。
lane等の用語についてはよく分かれていませんが、とりあえず実質的な処理はenqueueUpdateだけのようです。

ここは少し関数名だけで推測してしまっている部分もありますが、つまり、dispatchSetStateは、値の更新を適用(enqueue)するための関数と言えそうです。

それでは最後にmountStateに戻りましょう。

ふたたびmountState

コードは再掲しませんが、この関数は[hook.memoizedState, dispatch]をreturnしていましたね。
そしてhook.memolizedStateはinitialStateでした。
また、dispatchは値の更新を適用する関数でした。

なので、mountStateの返り値は[初期値, 更新用の関数]であると言えるでしょう。
これは初回時の[state, setState]を表しているので、ここで読解を終えることができそうです。

まとめ

つまり、useStateを初回に使った時は、以下のようにそれぞれの処理が動いていることになります。

useState ->
resolveDispatcher ->
ReactCurrentDispatcher.current ->
renderWithHooks ->
HooksDispatcherOnMount ->
mountState(とdispatchSetState)

複雑すぎ!!

所感

骨が折れる作業でしたが、OSSを読み進めていくということが初めてだったので、新鮮で、楽しかったです。
useStateの裏側を知ることが、実務のパフォーマンス向上につながることはないかもしれませんが、単純に面白かったので、よし。

参考記事

https://sbfl.net/blog/2019/02/09/react-hooks-usestate/
https://tnoyan.hatenablog.com/entry/2021/12/03/220348

Discussion