実装から学ぶ useMemo

はじめに

ReactHooks の1つである useMemo() を理解するため、 React の実装を読んだまとめです。
React の再レンダー時の余計な計算を減らすために useMemo を使うことがありますが、どんな仕組みで動いているのでしょうか。

React のバージョンは 19.0.0 です

コードを読む

まずは ReactHooks から見ていきましょう。

ReactHooks.js

export function useMemo<T>(
  create: () => T,
  deps: Array<mixed> | void | null,
): T {
  const dispatcher = resolveDispatcher();
  return dispatcher.useMemo(create, deps);
}

dispatcher の useMemo を呼び出しています。
マウント時には mountMemo が実行されます。

ReactFiberHooks.js

const HooksDispatcherOnMount: Dispatcher = {
  // ...
  useMemo: mountMemo,
  // ...
};

mountMemo では、関数を実行して、その結果と deps を一緒に hook.memoizedState にセットしています。

ReactFiberHooks.js

function mountMemo<T>(
  nextCreate: () => T,
  deps: Array<mixed> | void | null,
): T {
  const hook = mountWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  const nextValue = nextCreate();
  if (shouldDoubleInvokeUserFnsInHooksDEV) {
    setIsStrictModeForDevtools(true);
    nextCreate();
    setIsStrictModeForDevtools(false);
  }
  hook.memoizedState = [nextValue, nextDeps];
  return nextValue;
}

再レンダーされたときはどうでしょうか。

ReactFiberHooks.js

const HooksDispatcherOnRerender: Dispatcher = {
  // ...
  useMemo: updateMemo,
  // ...
};

updateMemo も見ていきましょう。

ReactFiberHooks.js

function updateMemo<T>(
  nextCreate: () => T,
  deps: Array<mixed> | void | null,
): T {
  const hook = updateWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  const prevState = hook.memoizedState;
  // Assume these are defined. If they're not, areHookInputsEqual will warn.
  if (nextDeps !== null) {
    const prevDeps: Array<mixed> | null = prevState[1];
    if (areHookInputsEqual(nextDeps, prevDeps)) {
      return prevState[0];
    }
  }
  // ...
}

まずは deps の変更を areHookInputsEqual でチェックしています。
念のためこの内容も見ておきましょう。

reactFiberHooks.js

function areHookInputsEqual(
  nextDeps: Array<mixed>,
  prevDeps: Array<mixed> | null,
): boolean {
  // ...
  for (let i = 0; i < prevDeps.length && i < nextDeps.length; i++) {
    if (is(nextDeps[i], prevDeps[i])) {
      continue;
    }
    return false;
  }
  return true;
}

依存配列の内容それぞれを is で比較しています。

objectIs.js

/**
 * inlined Object.is polyfill to avoid requiring consumers ship their own
 * https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Object/is
 */
function is(x: any, y: any) {
  return (
    (x === y && (x !== 0 || 1 / x === 1 / y)) || (x !== x && y !== y) // eslint-disable-line no-self-compare
  );
}

この is は Object.is のポリフィルです。
deps にオブジェクトや配列が渡された場合、内容が同じだったとしても参照が異なれば別の値として扱われます。

もう1度 updateMemo に戻りましょう。

ReactFiberHooks.js

function updateMemo<T>(
  nextCreate: () => T,
  deps: Array<mixed> | void | null,
): T {
  // ...
  if (nextDeps !== null) {
    const prevDeps: Array<mixed> | null = prevState[1];
    if (areHookInputsEqual(nextDeps, prevDeps)) {
      return prevState[0];
    }
  }
  const nextValue = nextCreate();
  if (shouldDoubleInvokeUserFnsInHooksDEV) {
    setIsStrictModeForDevtools(true);
    nextCreate();
    setIsStrictModeForDevtools(false);
  }
  hook.memoizedState = [nextValue, nextDeps];
  return nextValue;
}

areHookInputsEqual が true （つまり、 deps の値に変更がないとみなされた）の場合には prevState[0] = 前回の関数の計算結果が返されます。
それ以外の場合は nextCreate() で再度関数を実行し、マウント時と同様に hook.memoizedState に保存しています。

おわりに

なんだか難しそうな ReactHooks ですが、実装を読んでしまえば意外とシンプルですね！
他の Hooks にも言えることですが、 deps の比較が Object.is の浅い比較であることには注意が必要です。

やっていることがわかると利用イメージも湧きやすいので、今後もいろんな実装を読んでみようと思います。