useSelect: incrementally subscribe to stores when first selected from

packages/data/src/components/use-select/index.js

@@ -42,68 +42,116 @@ 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.
 			if ( lastMapResultValid && mapSelect === lastMapSelect ) {
 				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;