From 1e3e30dae2bcfbeb0abc686f2a37aec208eedb39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Chivot-Buhler?= <aurelien@arimaska.net>
Date: Tue, 8 Nov 2022 10:25:42 +0000
Subject: [PATCH] Fix useSyncExternalStore dropped update when state is
 dispatched in render phase (#25578)

Fix https://github.com/facebook/react/issues/25565
---
 .../src/ReactFiberHooks.new.js                |  2 +-
 .../src/ReactFiberHooks.old.js                |  2 +-
 .../__tests__/useSyncExternalStore-test.js    | 31 +++++++++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
index 1c1f2e6099373..7538bb9b3c217 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.new.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -1615,7 +1615,7 @@ function updateSyncExternalStore<T>(
       }
     }
   }
-  const prevSnapshot = hook.memoizedState;
+  const prevSnapshot = (currentHook || hook).memoizedState;
   const snapshotChanged = !is(prevSnapshot, nextSnapshot);
   if (snapshotChanged) {
     hook.memoizedState = nextSnapshot;
diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
index f9fac9358ad7b..e98d684a78e6f 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.old.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -1615,7 +1615,7 @@ function updateSyncExternalStore<T>(
       }
     }
   }
-  const prevSnapshot = hook.memoizedState;
+  const prevSnapshot = (currentHook || hook).memoizedState;
   const snapshotChanged = !is(prevSnapshot, nextSnapshot);
   if (snapshotChanged) {
     hook.memoizedState = nextSnapshot;
diff --git a/packages/react-reconciler/src/__tests__/useSyncExternalStore-test.js b/packages/react-reconciler/src/__tests__/useSyncExternalStore-test.js
index d2b838aebe0c0..804e40ca21afc 100644
--- a/packages/react-reconciler/src/__tests__/useSyncExternalStore-test.js
+++ b/packages/react-reconciler/src/__tests__/useSyncExternalStore-test.js
@@ -18,6 +18,7 @@ let useLayoutEffect;
 let forwardRef;
 let useImperativeHandle;
 let useRef;
+let useState;
 let startTransition;
 
 // This tests the native useSyncExternalStore implementation, not the shim.
@@ -36,6 +37,7 @@ describe('useSyncExternalStore', () => {
     useImperativeHandle = React.useImperativeHandle;
     forwardRef = React.forwardRef;
     useRef = React.useRef;
+    useState = React.useState;
     useSyncExternalStore = React.useSyncExternalStore;
     startTransition = React.startTransition;
 
@@ -173,4 +175,33 @@ describe('useSyncExternalStore', () => {
       });
     },
   );
+
+  test('next value is correctly cached when state is dispatched in render phase', async () => {
+    const store = createExternalStore('value:initial');
+
+    function App() {
+      const value = useSyncExternalStore(store.subscribe, store.getState);
+      const [sameValue, setSameValue] = useState(value);
+      if (value !== sameValue) setSameValue(value);
+      return <Text text={value} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    act(() => {
+      // Start a render that reads from the store and yields value
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded(['value:initial']);
+
+    await act(() => {
+      store.set('value:changed');
+    });
+    expect(Scheduler).toHaveYielded(['value:changed']);
+
+    // If cached value was updated, we expect a re-render
+    await act(() => {
+      store.set('value:initial');
+    });
+    expect(Scheduler).toHaveYielded(['value:initial']);
+  });
 });