memo.nnenn0.com

ReactのsetStateとFiberのつながりを理解するために実装を読んでみた

created_at
updated_at
tags
react

ReactのFiberとHookのつながりを理解するために実装を読んでみたの続き。前の記事では、countの値とsetCountがFiberからHook、queueとたどった先にあり、レンダー中のuseStateはcurrentlyRenderingFiberで自分のFiberを知ることを見た。今回も同じAppを題材にし、図では1つ目のCounterをCounter #1、2つ目をCounter #2と書く。

tsx
function Counter() {
  const [count, setCount] = useState(0);
  return <button onClick={() => setCount((n) => n + 1)}>{count}</button>;
}

function App() {
  return (
    <>
      <Counter />
      <Counter />
    </>
  );
}

2つ目のCounterのボタンを押してsetCount((n) => n + 1)を呼ぶとき、渡しているのは関数1つだけで、どちらのCounterかを示すものは渡していない。しかも、このとき動いているのはonClickの関数だけで、どのコンポーネントも呼ばれていないので、currentlyRenderingFiberはnullになっている。setCountは、レンダー中のuseStateのようにこの変数から自分のFiberを知ることはできない。知りたかったのは、setCountがどうやって2つ目のCounterに更新を届けているのかだ。

Fiberの中身は、前の記事と同じく、React v19.3.1でこのAppをNode.jsとjsdomに描画して、実際のオブジェクトを読んで確かめた。

マウント時のbind

setCountは、マウント中にuseStateを呼んだときにmountStateで作られる。型の注釈を省いて抜き出すと次のようになっている。

js
function mountState(initialState) {
  const hook = mountStateImpl(initialState);
  const queue = hook.queue;
  const dispatch = dispatchSetState.bind(
    null,
    currentlyRenderingFiber,
    queue,
  );
  queue.dispatch = dispatch;
  return [hook.memoizedState, dispatch];
}

ReactFiberHooks.js — mountState

mountStateImplは、Hookを作ってFiberのHook listにつなぎ、memoizedStateに初期値を入れ、queueを作る。その後のbindは、引数の一部を埋めた新しい関数を作るJavaScriptの機能だ。第1引数は新しい関数の中のthisで、dispatchSetStateはthisを使わないのでnullを渡している。第2引数以降は、dispatchSetState(fiber, queue, action)のfiberとqueueに埋め込まれる。

埋め込まれるのは、bindを呼んだ時点のcurrentlyRenderingFiberの中身だ。後でcurrentlyRenderingFiberが別のFiberやnullに変わっても、setCountに入ったFiberは変わらない。入っているのはFiberを複製したものではなく、ツリーにあるFiberそのものへの参照になる。

題材のAppをマウントすると、2つのCounterで次のようにsetCountが作られる。図では、Counter #1のqueueとsetCountをqueue #1、setCount #1と書き、Counter #2も同じように書く。

text
render Counter #1
  currentlyRenderingFiber = Counter #1
  useState(0) -> mountState
    create Hook, queue #1
    setCount #1 = dispatchSetState
     .bind(null, Counter #1, queue #1)
render Counter #2
  currentlyRenderingFiber = Counter #2
  useState(0) -> mountState
    create Hook, queue #2
    setCount #2 = dispatchSetState
     .bind(null, Counter #2, queue #2)

2つのCounterのsetCountは、実際に読んでも別の関数になっていた。2つ目のCounterのボタンでsetCount((n) => n + 1)を呼ぶと、dispatchSetState(Counter #2, queue #2, (n) => n + 1)が呼ばれる。ツリーを探して何番目のCounterかを調べなくても、setCountが最初から自分のFiberとqueueへの参照を持っている。

setCountはFiberのフィールドではなく、前の記事で見たようにqueueのdispatchに置かれている。FiberからはmemoizedState、Hook、queueとたどってsetCountに着き、setCountの側からは、bindで埋め込んだ参照でFiberとqueueを指している。

text
Counter #2 fiber
  memoizedState -> Hook
Hook
  queue         -> queue #2
queue #2
  dispatch      -> setCount #2
setCount #2 (bound)
  fiber         -> Counter #2 fiber
  queue         -> queue #2

setCountがFiberを使うところ

「マウント時のbind」の節で、setCountはFiberとqueueの両方への参照を持っていると書いた。Updateを登録するだけならqueueがあれば足りるのに、Fiberは何に使っているのか。setCountの中で動くdispatchSetStateInternalの最後を抜き出すと次のようになっている。

js
const root = enqueueConcurrentHookUpdate(fiber, queue, update, lane);
if (root !== null) {
  scheduleUpdateOnFiber(root, fiber, lane);
  // ...
}

ReactFiberHooks.js — dispatchSetStateInternal

fiberとqueueはbindで埋め込んだ参照で、updateは今回のUpdate、laneはこの更新の優先度だ。Fiberを使うのは次の2つになる。

では、このrootは何で、Fiberからどうやって探しているのか。

returnでたどるFiberRoot

dispatchSetStateInternalのrootは、createRootで作られるFiberRootというオブジェクトだ。前の記事で見たように、Fiberのツリーのいちばん上にはHostRootというFiberがあり、createRootはHostRootのFiberと一緒にFiberRootを作る（createFiberRoot）。FiberRootはFiberではなく、1つのReactのツリー全体を管理するオブジェクトで、コンテナのDOM（containerInfo）と、ツリーのいちばん上のFiber（current）を持っている（FiberRootNode）。

HostRootのstateNode（Fiberに対応する実体を指すフィールド）はFiberRootを指していて、実際に読んでもそのとおりにつながっていた。

text
FiberRoot
  containerInfo -> <div id="root">
  current       -> HostRoot
HostRoot
  stateNode     -> FiberRoot
  child         -> App

rootを探すのはgetRootForUpdatedFiberで、開発用のチェックを省いて抜き出すと次のようになっている。

js
let node = sourceFiber;
let parent = node.return;
while (parent !== null) {
  node = parent;
  parent = node.return;
}
return node.tag === HostRoot ? (node.stateNode as FiberRoot) : null;

ReactFiberConcurrentUpdates.js — getRootForUpdatedFiber

sourceFiberはsetCountが持っているFiberで、tagはFiberの種類を表すフィールドだ。returnを親がなくなるまでたどり、着いたFiberがHostRootなら、そのstateNodeのFiberRootを返す。2つ目のCounterなら、次のようにたどる。

text
Counter #2 fiber
  return    -> App
App
  return    -> HostRoot
HostRoot
  return    -> null (stop)
  stateNode -> FiberRoot

createRootを呼ぶたびにFiberRootとHostRootの組が1つずつできるので、どのツリーの更新なのかもこのたどり方で決まる。

コンポーネントがアンマウントされると、Reactは消したFiberのreturnをnullにしてツリーから切り離す（detachFiberMutation）。その後に残っていたsetCountが呼ばれても、たどった先がHostRootではないのでrootはnullになり、レンダーは予約されない。

FiberRootが見つかったときは、scheduleUpdateOnFiberがFiberRootのpendingLanes（そのツリーで保留中の更新のlane）にlaneを足し（markRootUpdated）、レンダーを予約する（scheduleUpdateOnFiber）。予約したレンダーは、FiberRootのcurrentが指すHostRootから始まる（prepareFreshStack）。そこからツリーを降りていき、2つ目のCounterを呼び直す。

currentとalternate

setCountが持っているのはマウント時のFiberだが、2つ目のCounterのボタンを押した後のレンダーで、Counterを呼ぶときに使われるFiberは、そのFiberとは限らない。

Reactは、1つの要素に対して、コミット済みのツリーにあるFiber（current）と、レンダー中に作るツリーにあるFiber（workInProgress）の2つを使う。2つはalternateで互いを指している。前の記事で見たrenderWithHooksのworkInProgressは、後者のFiberだ。レンダーではcurrentのFiberをもとにworkInProgressのFiberを用意する。alternateがまだなければ新しく作ってつなぎ、あれば使い回す（createWorkInProgress）。

レンダーを終えると、コミットでDOMを書き換えた後に、FiberRootのcurrentを付け替える。

js
root.current = finishedWork;

ReactFiberWorkLoop.js — flushMutationEffects

finishedWorkは、レンダーを終えたworkInProgressのツリーのHostRootだ。Fiberを1つずつ入れ替えるのではなく、この1か所を書き換えるだけで、workInProgressのツリーが新しいcurrentになる。Fiberには自分がcurrentかどうかを表すフィールドはなく、FiberRootのcurrentからたどって着く方がcurrentになる。

マウント時に作られた2つ目のCounterのFiberをA、1回目にボタンを押した後のレンダーで作られたもう1つをBとする。HostRootは、マウントの後にFiberRootのcurrentが指していた方を#1、もう1つを#2と呼ぶ。2つ目のCounterのボタンを2回押して、そのたびにFiberを読むと、次のように入れ替わっていた。

text
mount
  FiberRoot.current -> HostRoot #1
  Counter #2: A is current
1st click
  FiberRoot.current -> HostRoot #2
  Counter #2: B is current
2nd click
  FiberRoot.current -> HostRoot #1
  Counter #2: A is current again

マウントのレンダーではAだけが作られ、alternateはまだない。1回目のクリックの後のレンダーでは、AをもとにBを新しく作り、コミットでBがcurrentになる。2回目のクリックの後のレンダーでは、古くなったAがworkInProgressとして使い回され、コミットでAがcurrentに戻る。AのalternateはB、BのalternateはAのままで、3つ目のFiberはできない。

setCountは、ずっとマウント時のFiber（A）を持っている。Aが今のcurrentかどうかは、コミットのたびに変わる。

2つのFiberで共有するqueue

setCountが持っているFiberがcurrentでなくても、setCountが登録したUpdateは、次のレンダーのuseStateに届く。レンダーのたびにHookのオブジェクトは作り直されるが、queueだけは前のHookのものを引き継ぐからだ。

js
const newHook: Hook = {
  memoizedState: currentHook.memoizedState,
  // ...
  queue: currentHook.queue,
  // ...
};

ReactFiberHooks.js — updateWorkInProgressHook

currentHookは、currentのFiberにある同じ順番のHookだ。1回目に押した後に2つ目のCounterの2つのFiberを読むと、Hookは別のオブジェクトで、memoizedStateはBが1、Aが0だった。queueは同じオブジェクトだった。2回目に押した後も、memoizedStateはAが2、Bが1になり、queueは同じオブジェクトのままだった。

共有しているのはqueueというオブジェクトで、中のUpdateがずっと残るわけではない。レンダーでuseStateが呼ぶupdateReducerImplは、queue.pendingにたまったUpdateをHookの側へ移し、queue.pendingをnullに戻してから適用する（updateReducerImpl）。押した後に読んでも、queue.pendingはnullだった。

仮にqueueもコピーしていたら、setCountはbindで埋め込んだマウント時のqueueにUpdateを登録するのに、作り直したHookは別のqueueから読むことになり、更新が届かなくなる。

text
setCount #2
  bound queue -> queue #2
Hook in fiber A
  queue       -> queue #2
Hook in fiber B (if copied)
  queue       -> copy of queue #2

queueを2つのFiberで共有しているので、setCountがどちらのFiberを持っていても、Updateは同じqueueに届く。一方、Fiberのlanesは2つのFiberが別々に持っているので、Updateを登録するときに、Fiberとそのalternateの両方に印を付けている（enqueueUpdate）。

まとめ

setCountと、Fiber、queue、FiberRootは、次のようにつながっていた。

text
FiberRoot
  current -> HostRoot ... Counter
Counter fiber (A or B)
  memoizedState -> Hook
Hook
  memoizedState -> count
  queue         -> queue
queue
  dispatch      -> setCount
setCount (bound at mount)
  fiber         -> A
  queue         -> queue

setCountは、マウント時にbindで埋め込んだFiberとqueueへの参照を持っているので、レンダーの外から呼ばれても自分のFiberが分かる。FiberはreturnをたどってFiberRootを見つけ、レンダーを予約するのに使う。Fiberはコミットのたびにcurrentとalternateが入れ替わる。queueはその2つのFiberで共有しているので、setCountが登録した更新はどちらのFiberのレンダーにも届く。

FiberとHookのつながりを一通り追ってみて、今ならReact初心者だった頃の自分にReactを上手く教えてあげられる気がする。

予約したレンダーがHostRootからどうやって2つ目のCounterまでたどり着くのかは、続きのReactは更新のあるコンポーネントをどう見つけるのかで追った。

参考