💼

仕事が早く終わったので useCallback の実装を読んでみる

2024/09/12に公開

モチベーション

useCallback のドキュメントを読んでいて、以下の記載がありました。

すでに useMemo に詳しい場合、useCallback を次のように考えると役立つかもしれません。
https://ja.react.dev/reference/react/useCallback#how-is-usecallback-related-to-usememo

// Simplified implementation (inside React)
function useCallback(fn, dependencies) {
  return useMemo(() => fn, dependencies);
}

inside React とあったので useCallback の内部実装は useMemo を使ってるのかな、気になるな、見てみよう、と思ったのがモチベーションです。あと早めに仕事が終わったので、ゆったり react のぞいてみるか、というユルモチベです。

useMemo の挙動

丁寧に行きたいので、まずは useMemo からです。
useMemo は関数の呼び出し結果をキャッシュするフックで、挙動を整理すると以下になると思います。

1) コンポーネントのマウント時に初回は必ず実行され、その実行結果がキャッシュされる
2) コンポーネントが再レンダリングされた時に...
  2-A) 依存配列の値が変わっていなければキャッシュを返す
  2-B) 依存配列の値が変わっていれば再実行する

ただの感想ですが、実装を読んだところ、上記の挙動がそのままコードになっている印象でした。

コンポーネントのマウント時の処理は mountMemo という関数に実装されており、再レンダリング時の処理は updateMemo という関数に実装されています。

useMemo の実装（mountMemo と updateMemo）は ReactFiberHooks.js にあります。後述の useCallback の実装もここにあります。

mountMemo の実装

コンポーネントのマウント時に依存配列(deps)とメモ対象の実行結果(nextValue)をフックのステート(hook.memoizedState)に保った上で、メモ対象の実行結果(nextValue)返します。
shouldDoubleInvokeUserFnsInHooksDEV の分岐は、開発環境の Strict Mode を使用している場合のものだと思います。

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

実装内容を一言でまとめるなら、メモ対象の実行と実行結果のキャッシュです。

updateMemo の実装

マウント時に実行される mountMemo はシンプルでした。useMemo の本懐は再レンダリング時なので、useMemo の実装内容としては updateMemo の方が重要かもしれません。

コンポーネントの再レンダリング時に、マウント時の依存配列(prevDeps)と再レンダリング時の依存配列(nextDeps)を比較し、同一であれば、マウント時に保ったメモ対象の実行結果(prevState[0])を返します。

同一ではない場合、メモ対象を再実行し、フックのステート(hook.memoizedState)に保った上で、メモ対象の実行結果(nextValue)返します。

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];
    }
  }
  const nextValue = nextCreate();
  if (shouldDoubleInvokeUserFnsInHooksDEV) {
    setIsStrictModeForDevtools(true);
    nextCreate();
    setIsStrictModeForDevtools(false);
  }
  hook.memoizedState = [nextValue, nextDeps];
  return nextValue;
}

実装内容を一言でまとめるなら、キャッシュを返せるなら返して、返せないなら再実行している、です。

useCallback の挙動

次に本日メインテーマの useCallback を見ていきます。

useCallback は関数自体をキャッシュするフックで、挙動を整理すると以下になると思います。

1) コンポーネントのマウント時に関数自体がキャッシュされる
2) コンポーネントが再レンダリングされた時に...
  2-A) 依存配列の値が変わっていなければキャッシュを返す
  2-B) 依存配列の値が変わっていれば再実行する

useCallback の実装は useMemo とほぼ同じで、簡素に上記の挙動がそのままコードになっていました。構成も同じで、コンポーネントのマウント時は mountCallback に実装されており、再レンダリング時は updateCallback という関数に実装されています。

mountCallback の実装

以下の通り、ほぼ useMemo と同じで、違いはキャッシュ対象がコールバック関数(callback)である部分です。

function mountCallback<T>(callback: T, deps: Array<mixed> | void | null): T {
  const hook = mountWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  hook.memoizedState = [callback, nextDeps];
  return callback;
}

updateCallback の実装

こちらも以下の通り、ほぼ useMemo と同じで、違いはキャッシュ対象がコールバック関数(callback)である部分です。

function updateCallback<T>(callback: T, deps: Array<mixed> | void | null): T {
  const hook = updateWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  const prevState = hook.memoizedState;
  if (nextDeps !== null) {
    const prevDeps: Array<mixed> | null = prevState[1];
    if (areHookInputsEqual(nextDeps, prevDeps)) {
      return prevState[0];
    }
  }
  hook.memoizedState = [callback, nextDeps];
  return callback;
}

感想

useCallback と useMemo の実装差分はほぼありませんでした。また、気になっていた useCallback の内部実装は useMemo を使ってるのか、についてですが、使っていませんでした。

// Simplified implementation (inside React)
function useCallback(fn, dependencies) {
  return useMemo(() => fn, dependencies);
}

上記は概念的に useCallback を表現した記載でした。終わりです。