From 09f46e485246c40cc4dfa6cf42eda2a0008bd137 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Thu, 30 May 2024 12:45:30 -0400
Subject: [PATCH] Fix: Wrong dispatcher when replaying after suspend

When a component suspends with `use`, we switch to the "re-render"
dispatcher during the subsequent render attempt, so that we can reuse
the work from the initial attempt. However, once we run out of hooks
from the previous attempt, we should switch back to the regular
"update" dispatcher.

This is conceptually the same fix as the one introduced in #26232. That
fix only accounted for initial mount, but the useTransition regression
test added in the previous commit illustrates that we need to handle
updates, too.

The issue affects more than just useTransition but because most of the
behavior between the "re-render" and "update" dispatchers is the same
it's hard to contrive other scenarios in a test, which is probably why
it took so long for someone to notice.
---
 .../react-reconciler/src/ReactFiberHooks.js   | 53 ++++++++++++++-----
 .../src/__tests__/ReactUse-test.js            |  2 -
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js
index f69b8f1a2f3e3..aa01cdcad4234 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.js
@@ -1083,20 +1083,49 @@ function useThenable<T>(thenable: Thenable<T>): T {
     thenableState = createThenableState();
   }
   const result = trackUsedThenable(thenableState, thenable, index);
-  if (
-    currentlyRenderingFiber.alternate === null &&
-    (workInProgressHook === null
-      ? currentlyRenderingFiber.memoizedState === null
-      : workInProgressHook.next === null)
-  ) {
-    // Initial render, and either this is the first time the component is
-    // called, or there were no Hooks called after this use() the previous
-    // time (perhaps because it threw). Subsequent Hook calls should use the
-    // mount dispatcher.
+
+  // When something suspends with `use`, we replay the component with the
+  // "re-render" dispatcher instead of the "mount" or "update" dispatcher.
+  //
+  // But if there are additional hooks that occur after the `use` invocation
+  // that suspended, they wouldn't have been processed during the previous
+  // attempt. So after we invoke `use` again, we may need to switch from the
+  // "re-render" dispatcher back to the "mount" or "update" dispatcher. That's
+  // what the following logic accounts for.
+  //
+  // TODO: Theoretically this logic only needs to go into the rerender
+  // dispatcher. Could optimize, but probably not be worth it.
+
+  // This is the same logic as in updateWorkInProgressHook.
+  const workInProgressFiber = currentlyRenderingFiber;
+  const nextWorkInProgressHook =
+    workInProgressHook === null
+      ? // We're at the beginning of the list, so read from the first hook from
+        // the fiber.
+        workInProgressFiber.memoizedState
+      : workInProgressHook.next;
+
+  if (nextWorkInProgressHook !== null) {
+    // There are still hooks remaining from the previous attempt.
+  } else {
+    // There are no remaining hooks from the previous attempt. We're no longer
+    // in "re-render" mode. Switch to the normal mount or update dispatcher.
+    //
+    // This is the same as the logic in renderWithHooks, except we don't bother
+    // to track the hook types debug information in this case (sufficient to
+    // only do that when nothing suspends).
+    const currentFiber = workInProgressFiber.alternate;
     if (__DEV__) {
-      ReactSharedInternals.H = HooksDispatcherOnMountInDEV;
+      if (currentFiber !== null && currentFiber.memoizedState !== null) {
+        ReactSharedInternals.H = HooksDispatcherOnUpdateInDEV;
+      } else {
+        ReactSharedInternals.H = HooksDispatcherOnMountInDEV;
+      }
     } else {
-      ReactSharedInternals.H = HooksDispatcherOnMount;
+      ReactSharedInternals.H =
+        currentFiber === null || currentFiber.memoizedState === null
+          ? HooksDispatcherOnMount
+          : HooksDispatcherOnUpdate;
     }
   }
   return result;
diff --git a/packages/react-reconciler/src/__tests__/ReactUse-test.js b/packages/react-reconciler/src/__tests__/ReactUse-test.js
index 7374b8f4ffdf5..451912cd45478 100644
--- a/packages/react-reconciler/src/__tests__/ReactUse-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactUse-test.js
@@ -19,7 +19,6 @@ let useState;
 let useTransition;
 let useMemo;
 let useEffect;
-let useOptimistic;
 let Suspense;
 let startTransition;
 let pendingTextRequests;
@@ -43,7 +42,6 @@ describe('ReactUse', () => {
     useTransition = React.useTransition;
     useMemo = React.useMemo;
     useEffect = React.useEffect;
-    useOptimistic = React.useOptimistic;
     Suspense = React.Suspense;
     startTransition = React.startTransition;