Switch to mount dispatcher after use() when needed

packages/react-reconciler/src/ReactFiberHooks.js

@@ -658,10 +658,6 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
   // only get reset when the component either completes (finishRenderingHooks)
   // or unwinds (resetHooksOnUnwind).
   if (__DEV__) {
-    hookTypesDev =
-      current !== null
-        ? ((current._debugHookTypes: any): Array<HookType>)
-        : null;
     hookTypesUpdateIndexDev = -1;
     // Used for hot reloading:
     ignorePreviousDependencies =
@@ -696,8 +692,13 @@ function renderWithHooksAgain<Props, SecondArg>(
   let numberOfReRenders: number = 0;
   let children;
   do {
-    didScheduleRenderPhaseUpdateDuringThisPass = false;
+    if (didScheduleRenderPhaseUpdateDuringThisPass) {
+      // It's possible that a use() value depended on a state that was updated in
+      // this rerender, so we need to watch for different thenables this time.
+      thenableState = null;
+    }
     thenableIndexCounter = 0;
+    didScheduleRenderPhaseUpdateDuringThisPass = false;
 
     if (numberOfReRenders >= RE_RENDER_LIMIT) {
       throw new Error(
@@ -841,8 +842,7 @@ function updateWorkInProgressHook(): Hook {
   // This function is used both for updates and for re-renders triggered by a
   // render phase update. It assumes there is either a current hook we can
   // clone, or a work-in-progress hook from a previous render pass that we can
-  // use as a base. When we reach the end of the base list, we must switch to
-  // the dispatcher used for mounts.
+  // use as a base.
   let nextCurrentHook: null | Hook;
   if (currentHook === null) {
     const current = currentlyRenderingFiber.alternate;
@@ -876,16 +876,10 @@ function updateWorkInProgressHook(): Hook {
       if (currentFiber === null) {
         // This is the initial render. This branch is reached when the component
         // suspends, resumes, then renders an additional hook.
-        const newHook: Hook = {
-          memoizedState: null,
-
-          baseState: null,
-          baseQueue: null,
-          queue: null,
-
-          next: null,
-        };
-        nextCurrentHook = newHook;
+        // Should never be reached because we should switch to the mount dispatcher first.
+        throw new Error(
+          'Update hook called on initial render. This is likely a bug in React. Please file an issue.',
+        );
       } else {
         // This is an update. We should always have a current hook.
         throw new Error('Rendered more hooks than during the previous render.');
@@ -951,7 +945,24 @@ function use<T>(usable: Usable<T>): T {
       if (thenableState === null) {
         thenableState = createThenableState();
       }
-      return trackUsedThenable(thenableState, thenable, index);
+      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.
+        if (__DEV__) {
+          ReactCurrentDispatcher.current = HooksDispatcherOnMountInDEV;
+        } else {
+          ReactCurrentDispatcher.current = HooksDispatcherOnMount;
+        }
+      }
+      return result;
     } else if (
       usable.$$typeof === REACT_CONTEXT_TYPE ||
       usable.$$typeof === REACT_SERVER_CONTEXT_TYPE
@@ -1998,6 +2009,7 @@ function updateEffectImpl(
   const nextDeps = deps === undefined ? null : deps;
   let destroy = undefined;
 
+  // currentHook is null when rerendering after a render phase state update.
   if (currentHook !== null) {
     const prevEffect = currentHook.memoizedState;
     destroy = prevEffect.destroy;
@@ -2250,12 +2262,10 @@ 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 (prevState !== null) {
-    if (nextDeps !== null) {
-      const prevDeps: Array<mixed> | null = prevState[1];
-      if (areHookInputsEqual(nextDeps, prevDeps)) {
-        return prevState[0];
-      }
+  if (nextDeps !== null) {
+    const prevDeps: Array<mixed> | null = prevState[1];
+    if (areHookInputsEqual(nextDeps, prevDeps)) {
+      return prevState[0];
     }
   }
   hook.memoizedState = [callback, nextDeps];
@@ -2283,13 +2293,11 @@ function updateMemo<T>(
   const hook = updateWorkInProgressHook();
   const nextDeps = deps === undefined ? null : deps;
   const prevState = hook.memoizedState;
-  if (prevState !== null) {
-    // 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];
-      }
+  // 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];
     }
   }
   if (shouldDoubleInvokeUserFnsInHooksDEV) {

packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js

@@ -1076,7 +1076,7 @@ describe('ReactHooks', () => {
       expect(() => {
         ReactTestRenderer.create(<App />);
       }).toThrow(
-        'Should have a queue. This is likely a bug in React. Please file an issue.',
+        'Update hook called on initial render. This is likely a bug in React. Please file an issue.',
       );
     }).toErrorDev([
       'Do not call Hooks inside useEffect(...), useMemo(...), or other built-in Hooks',

packages/react-reconciler/src/__tests__/ReactThenable-test.js

@@ -5,6 +5,7 @@ let ReactNoop;
 let Scheduler;
 let act;
 let use;
+let useDebugValue;
 let useState;
 let useMemo;
 let Suspense;
@@ -21,6 +22,7 @@ describe('ReactThenable', () => {
     Scheduler = require('scheduler');
     act = require('jest-react').act;
     use = React.use;
+    useDebugValue = React.useDebugValue;
     useState = React.useState;
     useMemo = React.useMemo;
     Suspense = React.Suspense;
@@ -794,7 +796,7 @@ describe('ReactThenable', () => {
     expect(root).toMatchRenderedOutput('(empty)');
   });
 
-  test('when replaying a suspended component, reuses the hooks computed during the previous attempt', async () => {
+  test('when replaying a suspended component, reuses the hooks computed during the previous attempt (Memo)', async () => {
     function ExcitingText({text}) {
       // This computes the uppercased version of some text. Pretend it's an
       // expensive operation that we want to reuse.
@@ -846,6 +848,118 @@ describe('ReactThenable', () => {
     ]);
   });
 
+  test('when replaying a suspended component, reuses the hooks computed during the previous attempt (State)', async () => {
+    let _setFruit;
+    let _setVegetable;
+    function Kitchen() {
+      const [fruit, setFruit] = useState('apple');
+      _setFruit = setFruit;
+      const usedFruit = use(getAsyncText(fruit));
+      const [vegetable, setVegetable] = useState('carrot');
+      _setVegetable = setVegetable;
+      return <Text text={usedFruit + ' ' + vegetable} />;
+    }
+
+    // Initial render.
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(<Kitchen />);
+      });
+    });
+    expect(Scheduler).toHaveYielded(['Async text requested [apple]']);
+    expect(root).toMatchRenderedOutput(null);
+    await act(async () => {
+      resolveTextRequests('apple');
+    });
+    expect(Scheduler).toHaveYielded([
+      'Async text requested [apple]',
+      'apple carrot',
+    ]);
+    expect(root).toMatchRenderedOutput('apple carrot');
+
+    // Update the state variable after the use().
+    await act(async () => {
+      startTransition(() => {
+        _setVegetable('dill');
+      });
+    });
+    expect(Scheduler).toHaveYielded(['Async text requested [apple]']);
+    expect(root).toMatchRenderedOutput('apple carrot');
+    await act(async () => {
+      resolveTextRequests('apple');
+    });
+    expect(Scheduler).toHaveYielded([
+      'Async text requested [apple]',
+      'apple dill',
+    ]);
+    expect(root).toMatchRenderedOutput('apple dill');
+
+    // Update the state variable before the use(). The second state is maintained.
+    await act(async () => {
+      startTransition(() => {
+        _setFruit('banana');
+      });
+    });
+    expect(Scheduler).toHaveYielded(['Async text requested [banana]']);
+    expect(root).toMatchRenderedOutput('apple dill');
+    await act(async () => {
+      resolveTextRequests('banana');
+    });
+    expect(Scheduler).toHaveYielded([
+      'Async text requested [banana]',
+      'banana dill',
+    ]);
+    expect(root).toMatchRenderedOutput('banana dill');
+  });
+
+  test('when replaying a suspended component, reuses the hooks computed during the previous attempt (DebugValue+State)', async () => {
+    // Make sure we don't get a Hook mismatch warning on updates if there were non-stateful Hooks before the use().
+    let _setLawyer;
+    function Lexicon() {
+      useDebugValue(123);
+      const avocado = use(getAsyncText('aguacate'));
+      const [lawyer, setLawyer] = useState('abogado');
+      _setLawyer = setLawyer;
+      return <Text text={avocado + ' ' + lawyer} />;
+    }
+
+    // Initial render.
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(<Lexicon />);
+      });
+    });
+    expect(Scheduler).toHaveYielded(['Async text requested [aguacate]']);
+    expect(root).toMatchRenderedOutput(null);
+    await act(async () => {
+      resolveTextRequests('aguacate');
+    });
+    expect(Scheduler).toHaveYielded([
+      'Async text requested [aguacate]',
+      'aguacate abogado',
+    ]);
+    expect(root).toMatchRenderedOutput('aguacate abogado');
+
+    // Now update the state.
+    await act(async () => {
+      startTransition(() => {
+        _setLawyer('avocat');
+      });
+    });
+    expect(Scheduler).toHaveYielded(['Async text requested [aguacate]']);
+    expect(root).toMatchRenderedOutput('aguacate abogado');
+    await act(async () => {
+      resolveTextRequests('aguacate');
+    });
+    expect(Scheduler).toHaveYielded([
+      'Async text requested [aguacate]',
+      'aguacate avocat',
+    ]);
+    expect(root).toMatchRenderedOutput('aguacate avocat');
+  });
+
   // @gate enableUseHook
   test(
     'wrap an async function with useMemo to skip running the function ' +
@@ -1021,4 +1135,34 @@ describe('ReactThenable', () => {
     expect(Scheduler).toHaveYielded(['Async text requested [C]', 'C']);
     expect(root).toMatchRenderedOutput('ABC');
   });
+
+  // @gate enableUseHook
+  test('use() combined with render phase updates', async () => {
+    function Async() {
+      const a = use(Promise.resolve('A'));
+      const [count, setCount] = useState(0);
+      if (count === 0) {
+        setCount(1);
+      }
+      const usedCount = use(Promise.resolve(count));
+      return <Text text={a + usedCount} />;
+    }
+
+    function App() {
+      return (
+        <Suspense fallback={<Text text="Loading..." />}>
+          <Async />
+        </Suspense>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(<App />);
+      });
+    });
+    expect(Scheduler).toHaveYielded(['A1']);
+    expect(root).toMatchRenderedOutput('A1');
+  });
 });

scripts/error-codes/codes.json

@@ -451,5 +451,6 @@
   "463": "ReactDOMServer.renderToNodeStream(): The Node Stream API is not available in Bun. Use ReactDOMServer.renderToReadableStream() instead.",
   "464": "ReactDOMServer.renderToStaticNodeStream(): The Node Stream API is not available in Bun. Use ReactDOMServer.renderToReadableStream() instead.",
   "465": "enableFizzExternalRuntime without enableFloat is not supported. This should never appear in production, since it means you are using a misconfigured React bundle.",
-  "466": "Trying to call a function from \"use server\" but the callServer option was not implemented in your router runtime."
+  "466": "Trying to call a function from \"use server\" but the callServer option was not implemented in your router runtime.",
+  "467": "Update hook called on initial render. This is likely a bug in React. Please file an issue."
 }