diff --git a/packages/data/src/components/use-select/index.js b/packages/data/src/components/use-select/index.js index b41555a7da4495..fa7ffdaebb21b7 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,7 +128,17 @@ 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. @@ -103,7 +150,7 @@ function Store( registry, suspense ) { function getValue() { // Update the value in case it's been invalidated or `mapSelect` has changed. - updateValue( selectValue ); + updateValue(); return lastMapResult; } @@ -115,30 +162,13 @@ 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 8144611e3008b4..6967ae4e674923 100644 --- a/packages/data/src/components/use-select/test/index.js +++ b/packages/data/src/components/use-select/test/index.js @@ -21,10 +21,10 @@ import useSelect from '..'; jest.useRealTimers(); -function counterStore( initialCount = 0 ) { +function counterStore( initialCount = 0, step = 1 ) { return { reducer: ( state = initialCount, action ) => - action.type === 'INC' ? state + 1 : state, + action.type === 'INC' ? state + step : state, actions: { inc: () => ( { type: 'INC' } ), }, @@ -124,7 +124,7 @@ describe( 'useSelect', () => { ); expect( selectSpyFoo ).toHaveBeenCalledTimes( 2 ); - expect( selectSpyBar ).toHaveBeenCalledTimes( 2 ); + expect( selectSpyBar ).toHaveBeenCalledTimes( 1 ); expect( TestComponent ).toHaveBeenCalledTimes( 3 ); // Ensure expected state was rendered. @@ -195,6 +195,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