[useSES/extra] Reuse old selection if possible (#22307)

When you pass an unmemoized selector to useSyncExternalStoreExtra, we
have to reevaluate it on every render because we don't know whether
it depends on new props.

However, after reevalutating it, we should still compare the result
to the previous selection with `isEqual` and reuse the old one if it
hasn't changed.

Originally I did not implement this, because if the selector changes due
to new props or state, the component is going to have to re-render
anyway. However, it's still useful to return a memoized selection
when possible, because it may be the input to a downstream memoization.

In the test I wrote, the example I chose is selecting a list of
items from the store, and passing the list as a prop to a memoized
component. If the list prop is memoized, we can bail out.

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;
 }