Skip to content

Commit

Permalink
perf(hooks): Optimize useSyncExternalStoreWithSelector selector calls
Browse files Browse the repository at this point in the history
Add intelligent selector execution tracking to useSyncExternalStoreWithSelector
to avoid unnecessary re-renders and selector calls. The hook now tracks which
properties are accessed during selection and only re-runs the selector when
those specific properties change.

- Add Proxy-based property access tracking
- Track accessed properties to determine relevant state changes
- Skip selector execution when only irrelevant state changes
- Add tests verifying selector optimization behavior
- Preserve existing memoization and isEqual behavior

This improves performance by reducing unnecessary selector executions when
unrelated parts of the state change.
  • Loading branch information
Biki-das committed Jan 10, 2025
1 parent c80b336 commit 6da3ad1
Show file tree
Hide file tree
Showing 2 changed files with 300 additions and 17 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,217 @@
/**
* 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 ReactFeatureFlags;
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 {
// eslint-disable-next-line no-unused-vars
startTransition: _,
// eslint-disable-next-line no-unused-vars
useSyncExternalStore: __,
...otherExports
} = jest.requireActual('react');
return otherExports;
});
}

React = require('react');
ReactDOM = require('react-dom');
ReactDOMClient = require('react-dom/client');
ReactFeatureFlags = require('shared/ReactFeatureFlags');

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 expect(async () => {
await act(() => {
store.set({a: '2', b: '2'});
});
}).toWarnDev(
ReactFeatureFlags.enableUseRefAccessWarning
? ['Warning: App: Unsafe read of a mutable value during render.']
: [],
);

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);
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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 =
Expand Down

0 comments on commit 6da3ad1

Please sign in to comment.