diff --git a/packages/use-sync-external-store/src/__tests__/useSyncExternalStoreWithSelector-test.js b/packages/use-sync-external-store/src/__tests__/useSyncExternalStoreWithSelector-test.js new file mode 100644 index 0000000000000..c17e6e39ce124 --- /dev/null +++ b/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(); + }); + + 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(); + }); + + 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(); + }); + + expect(selectorFn).toHaveBeenCalledTimes(1); + + await act(() => { + store.set({a: '1', b: '3'}); + }); + + expect(selectorFn).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/use-sync-external-store/src/useSyncExternalStoreWithSelector.js b/packages/use-sync-external-store/src/useSyncExternalStoreWithSelector.js index a7be5c0fe8fda..3e03e95fabbee 100644 --- a/packages/use-sync-external-store/src/useSyncExternalStoreWithSelector.js +++ b/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( - subscribe: (() => void) => () => void, + subscribe: (onStoreChange: () => void) => () => void, getSnapshot: () => Snapshot, getServerSnapshot: void | null | (() => Snapshot), selector: (snapshot: Snapshot) => Selection, @@ -54,12 +54,44 @@ export function useSyncExternalStoreWithSelector( 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,8 +109,37 @@ export function useSyncExternalStoreWithSelector( } // 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. @@ -86,22 +147,27 @@ export function useSyncExternalStoreWithSelector( } // 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 =