From 397a93d85930d2efc186834f4d08f33f6bf1c9ff Mon Sep 17 00:00:00 2001
From: jerrydev0927 <jerrydev0927@gmail.com>
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 1c1f2e6099..7538bb9b3c 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 f9fac9358a..e98d684a78 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 d2b838aebe..804e40ca21 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']);
+  });
 });