From a6a27205a4ff9ec0450ca736156819231494d0c5 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Tue, 3 Oct 2023 15:01:40 -0400
Subject: [PATCH] Fix: Optimistic update does not get reset (#27453)

I found a bug where if an optimistic update causes a component to
rerender, and there are no other state updates during that render, React
bails out without applying the update.

Whenever a hook detects a change, we must mark the component as dirty to
prevent a bailout. We check for changes by comparing the new state to
`hook.memoizedState`. However, when implementing optimistic state
rebasing, I incorrectly reset `hook.memoizedState` to the incoming base
state, even though I only needed to reset `hook.baseState`. This was
just a mistake on my part.

This wasn't caught by the existing tests because usually when the
optimistic state changes, there's also some other state that marks the
component as dirty in the same render.

I fixed the bug and added a regression test.
---
 .../react-reconciler/src/ReactFiberHooks.js   | 12 ++--
 .../src/__tests__/ReactAsyncActions-test.js   | 63 +++++++++++++++++++
 2 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js
index 97f50d65f46c8..24746bde8d330 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.js
@@ -1817,9 +1817,9 @@ function updateOptimisticImpl<S, A>(
   // as an argument. It's called a passthrough because if there are no pending
   // updates, it will be returned as-is.
   //
-  // Reset the base state and memoized state to the passthrough. Future
-  // updates will be applied on top of this.
-  hook.baseState = hook.memoizedState = passthrough;
+  // Reset the base state to the passthrough. Future updates will be applied
+  // on top of this.
+  hook.baseState = passthrough;
 
   // If a reducer is not provided, default to the same one used by useState.
   const resolvedReducer: (S, A) => S =
@@ -1853,9 +1853,9 @@ function rerenderOptimistic<S, A>(
 
   // This is a mount. No updates to process.
 
-  // Reset the base state and memoized state to the passthrough. Future
-  // updates will be applied on top of this.
-  hook.baseState = hook.memoizedState = passthrough;
+  // Reset the base state to the passthrough. Future updates will be applied
+  // on top of this.
+  hook.baseState = passthrough;
   const dispatch = hook.queue.dispatch;
   return [passthrough, dispatch];
 }
diff --git a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js
index a333c4c486b4a..13143fa5ecf5b 100644
--- a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js
@@ -1112,4 +1112,67 @@ describe('ReactAsyncActions', () => {
       </>,
     );
   });
+
+  // @gate enableAsyncActions
+  test('useOptimistic can update repeatedly in the same async action', async () => {
+    let startTransition;
+    let setLoadingProgress;
+    let setText;
+    function App() {
+      const [, _startTransition] = useTransition();
+      const [text, _setText] = useState('A');
+      const [loadingProgress, _setLoadingProgress] = useOptimistic(0);
+      startTransition = _startTransition;
+      setText = _setText;
+      setLoadingProgress = _setLoadingProgress;
+
+      return (
+        <>
+          {loadingProgress !== 0 ? (
+            <div key="progress">
+              <Text text={`Loading... (${loadingProgress})`} />
+            </div>
+          ) : null}
+          <div key="real">
+            <Text text={text} />
+          </div>
+        </>
+      );
+    }
+
+    // Initial render
+    const root = ReactNoop.createRoot();
+    await act(() => root.render(<App text="A" />));
+    assertLog(['A']);
+    expect(root).toMatchRenderedOutput(<div>A</div>);
+
+    await act(async () => {
+      startTransition(async () => {
+        setLoadingProgress('25%');
+        await getText('Wait 1');
+        setLoadingProgress('75%');
+        await getText('Wait 2');
+        startTransition(() => setText('B'));
+      });
+    });
+    assertLog(['Loading... (25%)', 'A']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Loading... (25%)</div>
+        <div>A</div>
+      </>,
+    );
+
+    await act(() => resolveText('Wait 1'));
+    assertLog(['Loading... (75%)', 'A']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Loading... (75%)</div>
+        <div>A</div>
+      </>,
+    );
+
+    await act(() => resolveText('Wait 2'));
+    assertLog(['B']);
+  });
 });