[useSES/extra] Reuse old selection if possible

packages/use-sync-external-store/src/__tests__/useSyncExternalStoreShared-test.js

@@ -571,6 +571,8 @@ describe('Shared useSyncExternalStore behavior (shim and built-in)', () => {
   });
 
   describe('extra features implemented in user-space', () => {
+    // The selector implementation uses the lazy ref initialization pattern
+    // @gate !(enableUseRefAccessWarning && __DEV__)
     test('memoized selectors are only called once per update', async () => {
       const store = createExternalStore({a: 0, b: 0});
 
@@ -610,6 +612,8 @@ describe('Shared useSyncExternalStore behavior (shim and built-in)', () => {
       expect(root).toMatchRenderedOutput('A1');
     });
 
+    // The selector implementation uses the lazy ref initialization pattern
+    // @gate !(enableUseRefAccessWarning && __DEV__)
     test('Using isEqual to bailout', async () => {
       const store = createExternalStore({a: 0, b: 0});
 
@@ -666,4 +670,81 @@ describe('Shared useSyncExternalStore behavior (shim and built-in)', () => {
       expect(root).toMatchRenderedOutput('A1B1');
     });
   });
+
+  // The selector implementation uses the lazy ref initialization pattern
+  // @gate !(enableUseRefAccessWarning && __DEV__)
+  test('compares selection to rendered selection even if selector changes', async () => {
+    const store = createExternalStore({items: ['A', 'B']});
+
+    const shallowEqualArray = (a, b) => {
+      if (a.length !== b.length) {
+        return false;
+      }
+      for (let i = 0; i < a.length; i++) {
+        if (a[i] !== b[i]) {
+          return false;
+        }
+      }
+      return true;
+    };
+
+    const List = React.memo(({items}) => {
+      return (
+        <ul>
+          {items.map(text => (
+            <li key={text}>
+              <Text key={text} text={text} />
+            </li>
+          ))}
+        </ul>
+      );
+    });
+
+    function App({step}) {
+      const inlineSelector = state => {
+        Scheduler.unstable_yieldValue('Inline selector');
+        return [...state.items, 'C'];
+      };
+      const items = useSyncExternalStoreExtra(
+        store.subscribe,
+        store.getState,
+        inlineSelector,
+        shallowEqualArray,
+      );
+      return (
+        <>
+          <List items={items} />
+          <Text text={'Sibling: ' + step} />
+        </>
+      );
+    }
+
+    const root = createRoot();
+    await act(() => {
+      root.render(<App step={0} />);
+    });
+    expect(Scheduler).toHaveYielded([
+      'Inline selector',
+      'A',
+      'B',
+      'C',
+      'Sibling: 0',
+    ]);
+
+    await act(() => {
+      root.render(<App step={1} />);
+    });
+    expect(Scheduler).toHaveYielded([
+      // We had to call the selector again because it's not memoized
+      'Inline selector',
+
+      // But because the result was the same (according to isEqual) we can
+      // bail out of rendering the memoized list. These are skipped:
+      // 'A',
+      // 'B',
+      // 'C',
+
+      'Sibling: 1',
+    ]);
+  });
 });

packages/use-sync-external-store/src/useSyncExternalStoreExtra.js

@@ -13,7 +13,7 @@ import {useSyncExternalStore} from 'use-sync-external-store';
 
 // Intentionally not using named imports because Rollup uses dynamic
 // dispatch for CommonJS interop named imports.
-const {useMemo, useDebugValue} = React;
+const {useRef, useEffect, useMemo, useDebugValue} = React;
 
 // Same as useSyncExternalStore, but supports selector and isEqual arguments.
 export function useSyncExternalStoreExtra<Snapshot, Selection>(
@@ -22,6 +22,19 @@ export function useSyncExternalStoreExtra<Snapshot, Selection>(
   selector: (snapshot: Snapshot) => Selection,
   isEqual?: (a: Selection, b: Selection) => boolean,
 ): Selection {
+  // Use this to track the rendered snapshot.
+  const instRef = useRef(null);
+  let inst;
+  if (instRef.current === null) {
+    inst = {
+      hasValue: false,
+      value: (null: Selection | null),
+    };
+    instRef.current = inst;
+  } else {
+    inst = instRef.current;
+  }
+
   const getSnapshotWithMemoizedSelector = useMemo(() => {
     // Track the memoized state using closure variables that are local to this
     // memoized instance of a getSnapshot function. Intentionally not using a
@@ -38,6 +51,18 @@ export function useSyncExternalStoreExtra<Snapshot, Selection>(
         hasMemo = true;
         memoizedSnapshot = nextSnapshot;
         const nextSelection = selector(nextSnapshot);
+        if (isEqual !== undefined) {
+          // Even if the selector has changed, the currently rendered selection
+          // may be equal to the new selection. We should attempt to reuse the
+          // current value if possible, to preserve downstream memoizations.
+          if (inst.hasValue) {
+            const currentSelection = inst.value;
+            if (isEqual(currentSelection, nextSelection)) {
+              memoizedSelection = currentSelection;
+              return currentSelection;
+            }
+          }
+        }
         memoizedSelection = nextSelection;
         return nextSelection;
       }
@@ -67,10 +92,17 @@ export function useSyncExternalStoreExtra<Snapshot, Selection>(
       return nextSelection;
     };
   }, [getSnapshot, selector, isEqual]);
+
   const value = useSyncExternalStore(
     subscribe,
     getSnapshotWithMemoizedSelector,
   );
+
+  useEffect(() => {
+    inst.hasValue = true;
+    inst.value = value;
+  }, [value]);
+
   useDebugValue(value);
   return value;
 }