diff --git a/packages/data/src/components/use-select/index.js b/packages/data/src/components/use-select/index.js
index b41555a7da4495..e1082b50a54657 100644
--- a/packages/data/src/components/use-select/index.js
+++ b/packages/data/src/components/use-select/index.js
@@ -42,48 +42,85 @@ function Store( registry, suspense ) {
let lastMapResult;
let lastMapResultValid = false;
let lastIsAsync;
- let subscribe;
-
- const createSubscriber = ( stores ) => ( listener ) => {
- // Invalidate the value right after subscription was created. React will
- // call `getValue` after subscribing, to detect store updates that happened
- // in the interval between the `getValue` call during render and creating
- // the subscription, which is slightly delayed. We need to ensure that this
- // second `getValue` call will compute a fresh value.
- lastMapResultValid = false;
-
- const onStoreChange = () => {
- // Invalidate the value on store update, so that a fresh value is computed.
+ let subscriber;
+
+ const createSubscriber = ( stores ) => {
+ // The set of stores the `subscribe` function is supposed to subscribe to. Here it is
+ // initialized, and then the `updateStores` function can add new stores to it.
+ const activeStores = [ ...stores ];
+
+ // The `subscribe` function, which is passed to the `useSyncExternalStore` hook, could
+ // be called multiple times to establish multiple subscriptions. That's why we need to
+ // keep a set of active subscriptions;
+ const activeSubscriptions = new Set();
+
+ function subscribe( listener ) {
+ // Invalidate the value right after subscription was created. React will
+ // call `getValue` after subscribing, to detect store updates that happened
+ // in the interval between the `getValue` call during render and creating
+ // the subscription, which is slightly delayed. We need to ensure that this
+ // second `getValue` call will compute a fresh value.
lastMapResultValid = false;
- listener();
- };
- const onChange = () => {
- if ( lastIsAsync ) {
- renderQueue.add( queueContext, onStoreChange );
- } else {
- onStoreChange();
+ const onStoreChange = () => {
+ // Invalidate the value on store update, so that a fresh value is computed.
+ lastMapResultValid = false;
+ listener();
+ };
+
+ const onChange = () => {
+ if ( lastIsAsync ) {
+ renderQueue.add( queueContext, onStoreChange );
+ } else {
+ onStoreChange();
+ }
+ };
+
+ const unsubs = [];
+ function subscribeStore( storeName ) {
+ unsubs.push( registry.subscribe( onChange, storeName ) );
}
- };
- const unsubs = stores.map( ( storeName ) => {
- return registry.subscribe( onChange, storeName );
- } );
+ for ( const storeName of activeStores ) {
+ subscribeStore( storeName );
+ }
+
+ activeSubscriptions.add( subscribeStore );
+
+ return () => {
+ activeSubscriptions.delete( subscribeStore );
+
+ for ( const unsub of unsubs.values() ) {
+ // The return value of the subscribe function could be undefined if the store is a custom generic store.
+ unsub?.();
+ }
+ // Cancel existing store updates that were already scheduled.
+ renderQueue.cancel( queueContext );
+ };
+ }
+
+ // Check if `newStores` contains some stores we're not subscribed to yet, and add them.
+ function updateStores( newStores ) {
+ for ( const newStore of newStores ) {
+ if ( activeStores.includes( newStore ) ) {
+ continue;
+ }
+
+ // New `subscribe` calls will subscribe to `newStore`, too.
+ activeStores.push( newStore );
- return () => {
- // The return value of the subscribe function could be undefined if the store is a custom generic store.
- for ( const unsub of unsubs ) {
- unsub?.();
+ // Add `newStore` to existing subscriptions.
+ for ( const subscription of activeSubscriptions ) {
+ subscription( newStore );
+ }
}
- // Cancel existing store updates that were already scheduled.
- renderQueue.cancel( queueContext );
- };
- };
+ }
- return ( mapSelect, resubscribe, isAsync ) => {
- const selectValue = () => mapSelect( select, registry );
+ return { subscribe, updateStores };
+ };
- function updateValue( selectFromStore ) {
+ return ( mapSelect, isAsync ) => {
+ function updateValue() {
// If the last value is valid, and the `mapSelect` callback hasn't changed,
// then we can safely return the cached value. The value can change only on
// store update, and in that case value will be invalidated by the listener.
@@ -91,19 +128,30 @@ function Store( registry, suspense ) {
return lastMapResult;
}
- const mapResult = selectFromStore();
+ const listeningStores = { current: null };
+ const mapResult = registry.__unstableMarkListeningStores(
+ () => mapSelect( select, registry ),
+ listeningStores
+ );
+
+ if ( ! subscriber ) {
+ subscriber = createSubscriber( listeningStores.current );
+ } else {
+ subscriber.updateStores( listeningStores.current );
+ }
// If the new value is shallow-equal to the old one, keep the old one so
// that we don't trigger unwanted updates that do a `===` check.
if ( ! isShallowEqual( lastMapResult, mapResult ) ) {
lastMapResult = mapResult;
}
+ lastMapSelect = mapSelect;
lastMapResultValid = true;
}
function getValue() {
// Update the value in case it's been invalidated or `mapSelect` has changed.
- updateValue( selectValue );
+ updateValue();
return lastMapResult;
}
@@ -115,30 +163,12 @@ function Store( registry, suspense ) {
renderQueue.cancel( queueContext );
}
- // Either initialize the `subscribe` function, or create a new one if `mapSelect`
- // changed and has dependencies.
- // Usage without dependencies, `useSelect( ( s ) => { ... } )`, will subscribe
- // only once, at mount, and won't resubscibe even if `mapSelect` changes.
- if ( ! subscribe || ( resubscribe && mapSelect !== lastMapSelect ) ) {
- // Find out what stores the `mapSelect` callback is selecting from and
- // use that list to create subscriptions to specific stores.
- const listeningStores = { current: null };
- updateValue( () =>
- registry.__unstableMarkListeningStores(
- selectValue,
- listeningStores
- )
- );
- subscribe = createSubscriber( listeningStores.current );
- } else {
- updateValue( selectValue );
- }
+ updateValue();
lastIsAsync = isAsync;
- lastMapSelect = mapSelect;
// Return a pair of functions that can be passed to `useSyncExternalStore`.
- return { subscribe, getValue };
+ return { subscribe: subscriber.subscribe, getValue };
};
}
@@ -151,7 +181,7 @@ function useMappingSelect( suspense, mapSelect, deps ) {
const isAsync = useAsyncMode();
const store = useMemo( () => Store( registry, suspense ), [ registry ] );
const selector = useCallback( mapSelect, deps );
- const { subscribe, getValue } = store( selector, !! deps, isAsync );
+ const { subscribe, getValue } = store( selector, isAsync );
const result = useSyncExternalStore( subscribe, getValue, getValue );
useDebugValue( result );
return result;
diff --git a/packages/data/src/components/use-select/test/index.js b/packages/data/src/components/use-select/test/index.js
index f4b2ca83f24bf2..93ebf800be0a84 100644
--- a/packages/data/src/components/use-select/test/index.js
+++ b/packages/data/src/components/use-select/test/index.js
@@ -6,7 +6,7 @@ import { act, render, fireEvent, screen } from '@testing-library/react';
/**
* WordPress dependencies
*/
-import { Component, useState, useReducer } from '@wordpress/element';
+import { useLayoutEffect, useState, useReducer } from '@wordpress/element';
/**
* Internal dependencies
@@ -19,6 +19,19 @@ import {
} from '../../..';
import useSelect from '..';
+function counterStore( initialCount = 0, step = 1 ) {
+ return {
+ reducer: ( state = initialCount, action ) =>
+ action.type === 'INC' ? state + step : state,
+ actions: {
+ inc: () => ( { type: 'INC' } ),
+ },
+ selectors: {
+ get: ( state ) => state,
+ },
+ };
+}
+
describe( 'useSelect', () => {
let registry;
beforeEach( () => {
@@ -109,7 +122,7 @@ describe( 'useSelect', () => {
);
expect( selectSpyFoo ).toHaveBeenCalledTimes( 2 );
- expect( selectSpyBar ).toHaveBeenCalledTimes( 2 );
+ expect( selectSpyBar ).toHaveBeenCalledTimes( 1 );
expect( TestComponent ).toHaveBeenCalledTimes( 3 );
// Ensure expected state was rendered.
@@ -180,6 +193,81 @@ describe( 'useSelect', () => {
expect( Child ).toHaveBeenCalledTimes( 1 );
} );
+ it( 'incrementally subscribes to newly selected stores', () => {
+ registry.registerStore( 'store-main', counterStore() );
+ registry.registerStore( 'store-even', counterStore( 0, 2 ) );
+ registry.registerStore( 'store-odd', counterStore( 1, 2 ) );
+
+ const mapSelect = jest.fn( ( select ) => {
+ const first = select( 'store-main' ).get();
+ // select from other stores depending on whether main value is even or odd
+ const secondStore = first % 2 === 1 ? 'store-odd' : 'store-even';
+ const second = select( secondStore ).get();
+ return first + ':' + second;
+ } );
+
+ const TestComponent = jest.fn( () => {
+ const data = useSelect( mapSelect, [] );
+ return