From 2a77bb29c3f0592f21d6838e0877306a8a80a240 Mon Sep 17 00:00:00 2001
From: Ella <ella@vandurpe.com>
Date: Wed, 17 Jan 2024 23:43:34 +0100
Subject: [PATCH] Data: allow binding registry selector to multiple registries

---
 packages/data/src/factory.js                 | 17 ++++++++++-----
 packages/data/src/redux-store/index.js       |  5 +++++
 packages/data/src/test/registry-selectors.js | 23 ++++++++++++++++++--
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/packages/data/src/factory.js b/packages/data/src/factory.js
index 389c4a7e2b33ab..9bc14f46666c9b 100644
--- a/packages/data/src/factory.js
+++ b/packages/data/src/factory.js
@@ -39,17 +39,22 @@
  * @return {Function} Registry selector that can be registered with a store.
  */
 export function createRegistrySelector( registrySelector ) {
-	let selector;
-	let lastRegistry;
+	const selectorsByRegistry = new WeakMap();
 	// Create a selector function that is bound to the registry referenced by `selector.registry`
 	// and that has the same API as a regular selector. Binding it in such a way makes it
 	// possible to call the selector directly from another selector.
 	const wrappedSelector = ( ...args ) => {
-		if ( ! selector || lastRegistry !== wrappedSelector.registry ) {
-			selector = registrySelector( wrappedSelector.registry.select );
-			lastRegistry = wrappedSelector.registry;
+		// We want to make sure the cache persists even when new registry
+		// instances are created. For example patterns create their own editors
+		// with their own core/block-editor stores, so we should keep track of
+		// the cache for each registry instance.
+		if ( ! selectorsByRegistry.has( wrappedSelector.registry ) ) {
+			selectorsByRegistry.set(
+				wrappedSelector.registry,
+				registrySelector( wrappedSelector.registry.select )
+			);
 		}
-		return selector( ...args );
+		return selectorsByRegistry.get( wrappedSelector.registry )( ...args );
 	};
 
 	/**
diff --git a/packages/data/src/redux-store/index.js b/packages/data/src/redux-store/index.js
index e8394087215428..7fdc9331a2474b 100644
--- a/packages/data/src/redux-store/index.js
+++ b/packages/data/src/redux-store/index.js
@@ -237,6 +237,11 @@ export default function createReduxStore( key, options ) {
 				const boundSelector = ( ...args ) => {
 					args = normalize( selector, args );
 					const state = store.__unstableOriginalGetState();
+					// Before calling the selector, switch to the correct
+					// registry.
+					if ( selector.isRegistrySelector ) {
+						selector.registry = registry;
+					}
 					return selector( state.root, ...args );
 				};
 
diff --git a/packages/data/src/test/registry-selectors.js b/packages/data/src/test/registry-selectors.js
index bc567474541e24..edcadef8356c6b 100644
--- a/packages/data/src/test/registry-selectors.js
+++ b/packages/data/src/test/registry-selectors.js
@@ -78,8 +78,7 @@ describe( 'createRegistrySelector', () => {
 		expect( registry.select( uiStore ).getElementCount() ).toBe( 1 );
 	} );
 
-	// Even without createSelector, this fails in trunk.
-	it.skip( 'can bind one selector to multiple registries', () => {
+	it( 'can bind one selector to multiple registries (createRegistrySelector)', () => {
 		const registry1 = createRegistry();
 		const registry2 = createRegistry();
 
@@ -102,6 +101,26 @@ describe( 'createRegistrySelector', () => {
 		expect( registry2.select( uiStore ).getElementCount() ).toBe( 1 );
 	} );
 
+	it( 'can bind one selector to multiple registries (createRegistrySelector + createSelector)', () => {
+		const registry1 = createRegistry();
+		registry1.register( elementsStore );
+		registry1.register( uiStore );
+		registry1.dispatch( elementsStore ).add( 'Carbon' );
+
+		const registry2 = createRegistry();
+		registry2.register( elementsStore );
+		registry2.register( uiStore );
+		registry2.dispatch( elementsStore ).add( 'Helium' );
+
+		// Expects the `getFilteredElements` to be bound separately and independently to the two registries
+		expect( registry1.select( uiStore ).getFilteredElements() ).toEqual( [
+			'Carbon',
+		] );
+		expect( registry2.select( uiStore ).getFilteredElements() ).toEqual( [
+			'Helium',
+		] );
+	} );
+
 	it( 'can bind a memoized selector to a registry', () => {
 		const registry = createRegistry();
 		registry.register( elementsStore );