目的

React Hooksの中身の実装が気になって調べてみました。

実装を通じてReact Hooksの動作イメージを掴むことで、React Hooksの様々な制約（コールバックで呼べない、コンポーネント外で呼べない等）の根拠への理解など深まればいいかと思っています。

本記事のコードは全てReact 18.2.0で動作確認しています
なお筆者のReactスキルは趣味で1年触った程度で、React Hooksへの理解やReact実装の方式が誤っている可能性があります。予めご容赦ください

実装調査（修正：2022/12/06）

この章は読み飛ばしても問題ありません。Reactの詳細な内部実装に興味がない方は「ソースコード調査結果まとめ」または「ドキュメントによる情報」の項からご覧ください

試しにuseStateHookの実装を見てみました。
定義はReactHooks.js@L100にあり、

ReactHooks.js@L100

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

ReactHooks.js@L26

ReactHooks.js@L26

function resolveDispatcher() {
  const dispatcher = ReactCurrentDispatcher.current;
  if (__DEV__) {
    // ...
  }
  // ...
  return ((dispatcher: any): Dispatcher);
}

ReactCurrentDispatcher.js@L15

ReactCurrentDispatcher.js@L15

const ReactCurrentDispatcher = {
  current: (null: null | Dispatcher),
};

です。つまり、currentの値はデフォルトでnullのため、書き換え箇所を探す必要が出てきました。

ReactCurrentDispatcherの実体はreact-reconciler/src/ReactFiberHooks.jsの474行目で書き換えられています。（参考：https://sbfl.net/blog/2019/02/09/react-hooks-usestate/ ）

ReactFiberHooks.js@L474

export function renderWithHooks<Props, SecondArg>( // L418
  current: Fiber | null,
  // ...
): any {
  renderLanes = nextRenderLanes;
  currentlyRenderingFiber = workInProgress;
  if (__DEV__) {
    // ...
  } else {
    ReactCurrentDispatcher.current =
      current === null || current.memoizedState === null
        ? HooksDispatcherOnMount
        : HooksDispatcherOnUpdate; /// here!
  }
  //...
} // L545

ここでHooksDispatcherOnUpdateの中身は、同じファイルの2820行目に定義されています。

ReactFiberHooks.js@L2820

const HooksDispatcherOnUpdate: Dispatcher = {
  // ...
  useState: updateState,
  // ...
};

updateStateは1841行目に定義されていますが、その処理は単にupdateReducer（1058行目～）を呼び出しているだけです。
このupdateReducerの処理は非常に長いため全体はお示ししませんが、グローバルな値currentHookや動作中のHookオブジェクトupdateWorkInProgressHook()などを参照して処理を行っています。

ReactFiberHooks.js@L1058

function updateReducer<S, I, A>(
  reducer: (S, A) => S,
  initialArg: I,
  init?: I => S,
): [S, Dispatch<A>] {
  const hook = updateWorkInProgressHook();
  const queue = hook.queue;

  // ...

  const current: Hook = (currentHook: any); // L1074

  // ...

  if (baseQueue !== null) {
    // We have a queue to process.
    const first = baseQueue.next;
    let newState = current.baseState;

    let newBaseState = null;
    let newBaseQueueFirst = null;
    let newBaseQueueLast = null;
    let update = first;
    do {

        // ...

      if (shouldSkipUpdate) {
        // Priority is insufficient. Skip this update. If this is the first
        // skipped update, the previous update/state is the new base
        // update/state.

        // ...
      } else {
        // Process this update.
        const action = update.action;
        if (shouldDoubleInvokeUserFnsInHooksDEV) {
          reducer(newState, action);
        }
        if (update.hasEagerState) {
          // If this update is a state update (not a reducer) and was processed eagerly,
          // we can use the eagerly computed state
          newState = ((update.eagerState: any): S);
        } else {
          newState = reducer(newState, action);
        }
      }
      update = update.next;
    } while (update !== null && update !== first);

    // ...

    hook.memoizedState = newState;
    hook.baseState = newBaseState;
    hook.baseQueue = newBaseQueueLast;

    queue.lastRenderedState = newState;
  }

  // ...

  const dispatch: Dispatch<A> = (queue.dispatch: any);
  return [hook.memoizedState, dispatch];
}

ところで、ここまで見てきたファイルはReactFiberHooks.jsでした。ここで、Fiberというものについて簡単に説明すると、コンポーネントの（関数であれば）呼び出し等を管理するための描画単位です。¹

ソースコード調査結果まとめ

以上をまとめると、useStateHookは、Fiber（≒コンポーネント）ごとに保存されたデータを読み出し、その値と更新関数を作って返す実装になっているといえます。

ドキュメントによる情報

フックに関するよくある質問 - React: React はフック呼び出しとコンポーネントとをどのように関連付けているのですか？²
が見つかりました。
この記事によると、

それぞれのコンポーネントに関連付けられる形で、React 内に「メモリーセル」のリストが存在しています。それらは単に何らかのデータを保存できる JavaScript のオブジェクトです。あなたが useState() のようなフックを呼ぶと、フックは現在のセルの値を読み出し（あるいは初回レンダー時はセル内容を初期化し）、ポインタを次に進めます。これが複数の useState() の呼び出しが個別のローカル state を得る仕組みです。

とのことでした。
そのため、本記事では

コンポーネントに紐づいたデータ保存オブジェクト（メモリーセル）を作成・保持
データ保存オブジェクトを使ったHook（のようなもの）を実装
します。

Hook（らしきもの）を実装

本記事で実装した最終版のコードは https://github.com/stringthread/react-hooks-clone にて公開しています。

メモリーセル

各コンポーネントごとに、データ保存のためのオブジェクト群を作ります。とりあえず枠組みだけ作ってしまいましょう。

src/memoryCell.ts

import React from "react";
type CellValue = any;
type MemoryCell = {
    current: CellValue;
}; // 1つのデータを格納
type CellsForComponent = {}; // 1コンポーネント分のデータ群
type ComponentKey = str; // コンポーネントを識別するキーの型（後述）
type GlobalStore = Map<ComponentKey, CellsForComponent>; // プログラム全体のデータ群

MemoryCellのオブジェクトは何でも構いませんが、ここではMutableRefObjectに倣ってみました。

ドキュメント²によると、各コンポーネントの中では1つのデータを読み取ったらポインタを進めるそうです。
CellsForComponentに、この構造を追加しましょう。

src/memoryCell.ts

type CellsForComponent = {
    index: number; // データ群上の位置
    cells: MemoryCell[]; // データ群
};

合わせて、操作のための関数も追加しておきます。

src/memoryCell.ts

class CellsForComponent {
    index: number = 0; // データ群上の位置
    cells: MemoryCell[] = []; // データ群
    getAndNext(): MemoryCell {
        const prev_index = this.index;
        this.index++;
        if (this.index >= this.cells.length) {
            this.index=0;
        }
        return this.cells[prev_index];
    } // データ取得関数
    addCell(val: CellValue): MemoryCell {
        const newCell: MemoryCell = {
            current: val
        };
        this.cells.push(newCell);
        return newCell;
    } // 初期値を持ったMemoryCellを作ることで初期化
}

これでメモリーセルの基本は完成です。

Hookとコンポーネントを紐付け

Reactはレンダリングするとき、自分がどのコンポーネントであるか知っています。しかし、私たちが作る自作のuseStateは普通の関数ですから、どのコンポーネントから呼ばれたか知るのは簡単ではありません。（少なくとも私が調べた限り、関数コンポーネントの呼び出し時に新規コンポーネントか更新か判別する方法は見つかりませんでした）

基本Hookの自作とは言えなくなりますが、Reactが提供するuseStateで各コンポーネントのIDだけ管理し、それを使って判別することにします。

src/hookableFC.tsx

import { useState as useStateReactOfficial } from "react";
import { v4 as uuid } from "uuid";
const useComponentKey = (): ComponentKey => useStateReactOfficial(uuid())[0]; // 識別子を作るCustom Hook

次に、このuseComponentKeyHookをコンポーネントと関係付けます。

src/hookableFC.tsx

// 自作Hook用のKeyをpropsに持たせた型
type HookedFC<P extends object> = React.FC<
  P & { _component_key: ComponentKey }
>;

// コンポーネントに対してkeyを自動的に割り振るラッパ
// HookableFC<P>を受け取り、通常のReact.FC<P>を返す
const HookWrap = <P extends object>(Component: HookableFC<P>): React.FC<P> => (
    ( props: P ) => {
        const key = useComponentKey();
        return <Component {...props} _component_key={key} />;
    }
);

これで、HookedFCを使って定義されたコンポーネント内では、_component_keyを使ってコンポーネントを判別できるようになります。

実際には、React自体が内部で_component_keyのようなものを暗黙に持っていると思えばいいと思います
ただし、実コードで確認したわけではありません。誤っている可能性もありますがご容赦ください

ここで簡単に使い方の例を示します。

interface Prop {
    a: string;
}
const HockableSample: HookableFC<Prop> = (props) => (
    <div>{props._component_key}: {props.a}</div> // 内部で`_component_key`にアクセスできる
);
const Sample: React.FC<Prop> = HookWrap(HockableSample); // レンダリング時はこれを使う

// usage in rendering
const container = document.getElementById("root");
if (!container) throw Error("No container `#root` found");
const root = createRoot(container);
root.render(
  <>
    <Sample a={"1"} /> {/* `_component_key`は渡さなくていい */}
    <Sample a={"2"} />
  </>
);

Hook関数

ここまで、データ格納のメモリーセルと、コンポーネント識別機構を実装しました。最後に、useStateなどの実体にあたる関数を定義します。
いきなりuseStateで定義してもいいのですが、他のHookへの拡張を考慮して、まずはメモリーセルをそのまま返すuseCellを実装します。
このHookのすべきことを考えると、ドキュメント²で

あなたが useState() のようなフックを呼ぶと、フックは現在のセルの値を読み出し（あるいは初回レンダー時はセル内容を初期化し）、ポインタを次に進めます。

とある通り、初回レンダリングではセル初期化を、更新時にはセル内容の読み出しをそれぞれ行います。

つまり、自作Hookが呼ばれたとき、それが初回レンダリング中かどうか分からないといけません。この判定処理をHookWrap関数に追加しましょう。

src/hookableFC.tsx

  import { useEffect as useEffectReactOfficial } from "react";
+ // 初回レンダリングが終わったコンポーネントのKeyを格納
+ const componentsFinishedFirstRendering: Set<ComponentKey> = new Set();
  const HookWrap = <P extends object>(Component: HookableFC<P>): React.FC<P> => (
      ( props: P ) => {
          const key = useComponentKey(); // `react-hooks/rules-of-hooks`でLinterに怒られるが問題ない
          const rendered = <Component {...props} _component_key={key} />;
+         useEffectReactOfficial(()=>{
+             componentsFinishedFirstRendering.add(key); // レンダリング後にKeyを追加
+         },[]);
          return rendered;
      }
  );

これで、KeyがcomponentsFinishedFirstRenderingにあるかどうかで、初回レンダリングかどうか判定できるようになりました。

いよいよ、Hook本体を実装します。

src/hooks.ts

const emptyInitializationSymbol = Symbol();
const useCell = (key: ComponentKey): MemoryCell => {
    if (!store.has(key)) store.set(key, new CellsForComponent());
    const cells = store.get(key) as CellsForComponent; // 簡単のため、例外チェックを省略
    if (componentsFinishedFirstRendering.has(key)) {
        return cells.getAndNext(); // 2回目以降なら単にセルを読み出し
    }
    return cells.addCell(emptyInitializationSymbol); // 初回はセルを新規作成で初期化
};

これで本質的な部分は完成ですが、試しにuseStateHookも作ってみましょう。

src/hooks.ts

import _ from 'lodash';

const useState = <T,>(initial: T, key: ComponentKey): [T, (val: T) => void] => {
    const cell = useCell(key);
    // 初回は引数で値初期化
    if(cell.current === emptyInitializationSymbol) {
        cell.current = initial;
    }
    // setter関数ではcell.currentを書き換え
    const setter = (val: T): void => {
        cell.current = val;
    };
    // 値はcloneDeepして返す
    return [_.cloneDeep(cell.current), setter];
};

以上でメモリーセルの読み書きができました。使ってみましょう。

// src/sample/sampleComponent.tsx

interface SampleProps {
  initial: number;
}
const HockableSampleComponent: HookableFC<SampleProps> = (props) => {
  const [count, setCount] = useState(props.initial, props._component_key);
  return (
    <>
      <div>{count}</div>
      <button onClick={() => {
        setCount(count + 1);
        console.log(`called: ${store.get(props._component_key)?.cells.map(v=>v.current)}`); // メモリーセルの中身
      }}>+1</button>
    </>
  );
};
const Sample: React.FC<Prop> = HookWrap(HockableSampleComponent);

// src/sample/index.tsx

const container = document.getElementById("root");
if (!container) throw Error("No container `#root` found");
const root = createRoot(container);
root.render(
  <>
    <Sample initial={1} /> {/* `_component_key`は渡さなくていい */}
  </>
);

実際に動かしてみると、メモリーセルの値は変わるものの再レンダリングが行われません。

実は（というほどでもありませんが）React公式が提供するuseStateHookは、状態の書き換えに加えて再レンダリング作用も持っています。

先ほど実装したuseStateHookにも同様の処理を付け足してみましょう。³

src/hooks.ts

+ const useForceUpdate = () => {
+     const [_, setDummyState] = useStateReactOfficial(0); 
+     return () => setDummyState(e => e + 1); 
+ };
  const useState = <T,>(initial: T, key: ComponentKey): [T, (val: T) => void] => {
      const cell = useCell(key);
      // 初回は引数で値初期化
      if(cell.current === emptyInitializationSymbol) {
          cell.current = initial;
      }
+     const forceUpdate = useForceUpdate();
      // setter関数ではcell.currentを書き換え
      const setter = (val: T): void => {
          cell.current = val;
+         forceUpdate();
      };
      // 値はcloneDeepして返す
      return [_.cloneDeep(cell.current), setter];
  };

これで自作のuseStateHookが正常に動作するようになりました。

おわりに

Reactのレンダリングプロセスに関連する部分は一部React公式のHooksに頼ってしまいましたが、React Hooksの動作について何となくイメージできたのではないでしょうか。（といっても、ドキュメント²の情報をそのまま実装しただけですが……）

React Hooksの動作イメージとして「コンポーネントごとにメモリーセルの配列があり、それを順に読み出している」という点が分かったことは1つ大きく、これがReact Hooksの使用方法に関する制約の元になっていることを改めて体感できました。

他のHookも大抵はuseCellを元に実装できるはずです。お暇な方は適当に実装して、下にあるGitHubリポジトリにPR投げてください。
ちなみに筆者はuseEffect系の実装について何の着想も持っていないので、もし方策を思いついた方はコメント頂けると大変嬉しいです。

コード全体

https://github.com/stringthread/react-hooks-clone にて公開しています。

React Fiber Architecture: README.md
https://github.com/acdlite/react-fiber-architecture ↩
フックに関するよくある質問 - React: React はフック呼び出しとコンポーネントとをどのように関連付けているのですか？ https://ja.reactjs.org/docs/hooks-faq.html#how-does-react-associate-hook-calls-with-components ↩ ↩² ↩³ ↩⁴
React コンポーネントを強制的に再レンダリングする方法 https://dev-k.hatenablog.com/entry/how-to-force-re-rendering-of-react-components-dev-k ↩