perf(hooks): Optimize useSyncExternalStoreWithSelector selector calls

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

@@ -0,0 +1,206 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+let useSyncExternalStoreWithSelector;
+let React;
+let ReactDOM;
+let ReactDOMClient;
+let act;
+
+describe('useSyncExternalStoreWithSelector', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    if (gate(flags => flags.enableUseSyncExternalStoreShim)) {
+      // Remove useSyncExternalStore from the React imports so that we use the
+      // shim instead. Also removing startTransition, since we use that to
+      // detect outdated 18 alphas that don't yet include useSyncExternalStore.
+      //
+      // Longer term, we'll probably test this branch using an actual build
+      // of React 17.
+      jest.mock('react', () => {
+        const {
+          startTransition: _,
+          useSyncExternalStore: __,
+          ...otherExports
+        } = jest.requireActual('react');
+        return otherExports;
+      });
+    }
+
+    React = require('react');
+    ReactDOM = require('react-dom');
+    ReactDOMClient = require('react-dom/client');
+    const internalAct = require('internal-test-utils').act;
+
+    // The internal act implementation doesn't batch updates by default, since
+    // it's mostly used to test concurrent mode. But since these tests run
+    // in both concurrent and legacy mode, I'm adding batching here.
+    act = cb => internalAct(() => ReactDOM.unstable_batchedUpdates(cb));
+
+    if (gate(flags => flags.source)) {
+      // The `shim/with-selector` module composes the main
+      // `use-sync-external-store` entrypoint. In the compiled artifacts, this
+      // is resolved to the `shim` implementation by our build config, but when
+      // running the tests against the source files, we need to tell Jest how to
+      // resolve it. Because this is a source module, this mock has no affect on
+      // the build tests.
+      jest.mock('use-sync-external-store/src/useSyncExternalStore', () =>
+        jest.requireActual('use-sync-external-store/shim'),
+      );
+    }
+    useSyncExternalStoreWithSelector =
+      require('use-sync-external-store/shim/with-selector').useSyncExternalStoreWithSelector;
+  });
+
+  function createRoot(container) {
+    // This wrapper function exists so we can test both legacy roots and
+    // concurrent roots.
+    if (gate(flags => !flags.enableUseSyncExternalStoreShim)) {
+      // The native implementation only exists in 18+, so we test using
+      // concurrent mode.
+      return ReactDOMClient.createRoot(container);
+    } else {
+      // For legacy mode, use ReactDOM.createRoot instead of ReactDOM.render
+      const root = ReactDOMClient.createRoot(container);
+      return {
+        render(children) {
+          root.render(children);
+        },
+      };
+    }
+  }
+
+  function createExternalStore(initialState) {
+    const listeners = new Set();
+    let currentState = initialState;
+    return {
+      set(text) {
+        currentState = text;
+        ReactDOM.unstable_batchedUpdates(() => {
+          listeners.forEach(listener => listener());
+        });
+      },
+      subscribe(listener) {
+        listeners.add(listener);
+        return () => listeners.delete(listener);
+      },
+      getState() {
+        return currentState;
+      },
+      getSubscriberCount() {
+        return listeners.size;
+      },
+    };
+  }
+
+  test('should call selector on change accessible segment', async () => {
+    const store = createExternalStore({a: '1', b: '2'});
+
+    const selectorFn = jest.fn();
+    const selector = state => {
+      selectorFn();
+      return state.a;
+    };
+
+    function App() {
+      const data = useSyncExternalStoreWithSelector(
+        store.subscribe,
+        store.getState,
+        null,
+        selector,
+      );
+      return <>{data}</>;
+    }
+
+    const container = document.createElement('div');
+    const root = createRoot(container);
+    await act(() => {
+      root.render(<App />);
+    });
+
+    expect(selectorFn).toHaveBeenCalledTimes(1);
+
+    await act(() => {
+      store.set({a: '2', b: '2'});
+    });
+
+    expect(selectorFn).toHaveBeenCalledTimes(2);
+  });
+
+  test('should not call selector if nothing changed', async () => {
+    const store = createExternalStore({a: '1', b: '2'});
+
+    const selectorFn = jest.fn();
+    const selector = state => {
+      selectorFn();
+      return state.a;
+    };
+
+    function App() {
+      const data = useSyncExternalStoreWithSelector(
+        store.subscribe,
+        store.getState,
+        null,
+        selector,
+      );
+      return <>{data}</>;
+    }
+
+    const container = document.createElement('div');
+    const root = createRoot(container);
+    await act(() => {
+      root.render(<App />);
+    });
+
+    expect(selectorFn).toHaveBeenCalledTimes(1);
+
+    await act(() => {
+      store.set({a: '1', b: '2'});
+    });
+
+    expect(selectorFn).toHaveBeenCalledTimes(1);
+  });
+
+  test('should not call selector on change not accessible segment', async () => {
+    const store = createExternalStore({a: '1', b: '2'});
+
+    const selectorFn = jest.fn();
+    const selector = state => {
+      selectorFn();
+      return state.a;
+    };
+
+    function App() {
+      const data = useSyncExternalStoreWithSelector(
+        store.subscribe,
+        store.getState,
+        null,
+        selector,
+      );
+      return <>{data}</>;
+    }
+
+    const container = document.createElement('div');
+    const root = createRoot(container);
+    await act(() => {
+      root.render(<App />);
+    });
+
+    expect(selectorFn).toHaveBeenCalledTimes(1);
+
+    await act(() => {
+      store.set({a: '1', b: '3'});
+    });
+
+    expect(selectorFn).toHaveBeenCalledTimes(1);
+  });
+});

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

@@ -17,7 +17,7 @@ const {useRef, useEffect, useMemo, useDebugValue} = React;
 
 // Same as useSyncExternalStore, but supports selector and isEqual arguments.
 export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
-  subscribe: (() => void) => () => void,
+  subscribe: (onStoreChange: () => void) => () => void,
   getSnapshot: () => Snapshot,
   getServerSnapshot: void | null | (() => Snapshot),
   selector: (snapshot: Snapshot) => Selection,
@@ -54,12 +54,44 @@ export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
     let hasMemo = false;
     let memoizedSnapshot;
     let memoizedSelection: Selection;
+    let lastUsedProps: string[] = [];
+    let hasAccessed = false;
+    const accessedProps: string[] = [];
+
     const memoizedSelector = (nextSnapshot: Snapshot) => {
+      const getProxy = (): Snapshot => {
+        if (
+          !(typeof nextSnapshot === 'object') ||
+          typeof Proxy === 'undefined'
+        ) {
+          return nextSnapshot;
+        }
+
+        const handler = {
+          get: (target: Snapshot, prop: string, receiver: any) => {
+            const propertyName = prop.toString();
+
+            if (accessedProps.indexOf(propertyName) === -1) {
+              accessedProps.push(propertyName);
+            }
+
+            const value = Reflect.get(target, prop, receiver);
+
+            return value;
+          },
+        };
+
+        return (new Proxy(nextSnapshot, handler): any);
+      };
+
       if (!hasMemo) {
         // The first time the hook is called, there is no memoized result.
         hasMemo = true;
         memoizedSnapshot = nextSnapshot;
-        const nextSelection = selector(nextSnapshot);
+        const nextSelection = selector(getProxy());
+        lastUsedProps = accessedProps;
+        hasAccessed = true;
+
         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
@@ -77,31 +109,65 @@ export function useSyncExternalStoreWithSelector<Snapshot, Selection>(
       }
 
       // We may be able to reuse the previous invocation's result.
-      const prevSnapshot: Snapshot = (memoizedSnapshot: any);
-      const prevSelection: Selection = (memoizedSelection: any);
+      const prevSnapshot = memoizedSnapshot;
+      const prevSelection = memoizedSelection;
+
+      const getChangedSegments = (): string[] | void => {
+        if (
+          prevSnapshot === undefined ||
+          !hasAccessed ||
+          lastUsedProps.length === 0
+        ) {
+          return undefined;
+        }
+
+        const result: string[] = [];
+
+        if (
+          nextSnapshot !== null &&
+          typeof nextSnapshot === 'object' &&
+          prevSnapshot !== null &&
+          typeof prevSnapshot === 'object'
+        ) {
+          for (let i = 0; i < lastUsedProps.length; i++) {
+            const segmentName = lastUsedProps[i];
+
+            if (nextSnapshot[segmentName] !== prevSnapshot[segmentName]) {
+              result.push(segmentName);
+            }
+          }
+        }
+
+        return result;
+      };
 
       if (is(prevSnapshot, nextSnapshot)) {
         // The snapshot is the same as last time. Reuse the previous selection.
         return prevSelection;
       }
 
       // The snapshot has changed, so we need to compute a new selection.
-      const nextSelection = selector(nextSnapshot);
-
-      // If a custom isEqual function is provided, use that to check if the data
-      // has changed. If it hasn't, return the previous selection. That signals
-      // to React that the selections are conceptually equal, and we can bail
-      // out of rendering.
-      if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
-        // The snapshot still has changed, so make sure to update to not keep
-        // old references alive
+
+      const changedSegments = getChangedSegments();
+      if (changedSegments === undefined || changedSegments.length > 0) {
+        const nextSelection = selector(getProxy());
+        lastUsedProps = accessedProps;
+        hasAccessed = true;
+
+        // If a custom isEqual function is provided, use that to check if the data
+        // has changed. If it hasn't, return the previous selection. That signals
+        // to React that the selections are conceptually equal, and we can bail
+        // out of rendering.
+        if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
+          return prevSelection;
+        }
+
         memoizedSnapshot = nextSnapshot;
+        memoizedSelection = nextSelection;
+        return nextSelection;
+      } else {
         return prevSelection;
       }
-
-      memoizedSnapshot = nextSnapshot;
-      memoizedSelection = nextSelection;
-      return nextSelection;
     };
     // Assigning this to a constant so that Flow knows it can't change.
     const maybeGetServerSnapshot =