From f2c381131fb58c107e6153255bc8d1d6340db506 Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Mon, 4 Oct 2021 23:39:51 +0900
Subject: [PATCH] fix: useSyncExternalStoreExtra (#22500)

* move setting memoizedSnapshot

* Revert "move setting memoizedSnapshot"

This reverts commit f01320689cebfcbc4f3a53208f879ed4a8d6613d.

* add failed tests

* Revert "Revert "move setting memoizedSnapshot""

This reverts commit cb43c4fdc6b0dcab3480f27d6fbbb3137dbc47bb.
---
 .../useSyncExternalStoreShared-test.js        | 89 +++++++++++++++++++
 .../src/useSyncExternalStoreExtra.js          |  2 +-
 2 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/packages/use-sync-external-store/src/__tests__/useSyncExternalStoreShared-test.js b/packages/use-sync-external-store/src/__tests__/useSyncExternalStoreShared-test.js
index 6536f65b7f2a2..ab8143f24034c 100644
--- a/packages/use-sync-external-store/src/__tests__/useSyncExternalStoreShared-test.js
+++ b/packages/use-sync-external-store/src/__tests__/useSyncExternalStoreShared-test.js
@@ -818,4 +818,93 @@ describe('Shared useSyncExternalStore behavior (shim and built-in)', () => {
       'Sibling: 1',
     ]);
   });
+
+  describe('selector and isEqual error handling in extra', () => {
+    let ErrorBoundary;
+    beforeAll(() => {
+      spyOnDev(console, 'warn');
+      ErrorBoundary = class extends React.Component {
+        state = {error: null};
+        static getDerivedStateFromError(error) {
+          return {error};
+        }
+        render() {
+          if (this.state.error) {
+            return <Text text={this.state.error.message} />;
+          }
+          return this.props.children;
+        }
+      };
+    });
+
+    it('selector can throw on update', async () => {
+      const store = createExternalStore({a: 'a'});
+      const selector = state => state.a.toUpperCase();
+
+      function App() {
+        const a = useSyncExternalStoreExtra(
+          store.subscribe,
+          store.getState,
+          null,
+          selector,
+        );
+        return <Text text={a} />;
+      }
+
+      const container = document.createElement('div');
+      const root = createRoot(container);
+      await act(() =>
+        root.render(
+          <ErrorBoundary>
+            <App />
+          </ErrorBoundary>,
+        ),
+      );
+
+      expect(container.textContent).toEqual('A');
+
+      await act(() => {
+        store.set({});
+      });
+      expect(container.textContent).toEqual(
+        "Cannot read property 'toUpperCase' of undefined",
+      );
+    });
+
+    it('isEqual can throw on update', async () => {
+      const store = createExternalStore({a: 'A'});
+      const selector = state => state.a;
+      const isEqual = (left, right) => left.a.trim() === right.a.trim();
+
+      function App() {
+        const a = useSyncExternalStoreExtra(
+          store.subscribe,
+          store.getState,
+          null,
+          selector,
+          isEqual,
+        );
+        return <Text text={a} />;
+      }
+
+      const container = document.createElement('div');
+      const root = createRoot(container);
+      await act(() =>
+        root.render(
+          <ErrorBoundary>
+            <App />
+          </ErrorBoundary>,
+        ),
+      );
+
+      expect(container.textContent).toEqual('A');
+
+      await act(() => {
+        store.set({});
+      });
+      expect(container.textContent).toEqual(
+        "Cannot read property 'trim' of undefined",
+      );
+    });
+  });
 });
diff --git a/packages/use-sync-external-store/src/useSyncExternalStoreExtra.js b/packages/use-sync-external-store/src/useSyncExternalStoreExtra.js
index f4a1885aec5b9..aa4957b534753 100644
--- a/packages/use-sync-external-store/src/useSyncExternalStoreExtra.js
+++ b/packages/use-sync-external-store/src/useSyncExternalStoreExtra.js
@@ -76,7 +76,6 @@ export function useSyncExternalStoreExtra<Snapshot, Selection>(
       }
 
       // The snapshot has changed, so we need to compute a new selection.
-      memoizedSnapshot = nextSnapshot;
       const nextSelection = selector(nextSnapshot);
 
       // If a custom isEqual function is provided, use that to check if the data
@@ -87,6 +86,7 @@ export function useSyncExternalStoreExtra<Snapshot, Selection>(
         return prevSelection;
       }
 
+      memoizedSnapshot = nextSnapshot;
       memoizedSelection = nextSelection;
       return nextSelection;
     };