Switch to mount dispatcher after use() when needed (#26232)

sophiebits · web-flow · commit a8f971b7a669 · 2023-02-24T12:06:27.000-08:00
When resuming a suspended render, there may be more Hooks to be called that weren't seen the previous time through. Make sure to switch to the mount dispatcher when calling use() if the next Hook call should be treated as a mount. Fixes #25964.
diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/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) {
diff --git a/packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js b/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',
diff --git a/packages/react-reconciler/src/__tests__/ReactThenable-test.js b/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');
+  });
 });
diff --git a/scripts/error-codes/codes.json b/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."
 }
