Support registry inheritence with atomic stores (#27162)

packages/data/src/atomic-store/index.js

@@ -6,7 +6,11 @@ import { mapValues } from 'lodash';
 /**
  * WordPress dependencies
  */
-import { createDerivedAtom } from '@wordpress/stan';
+import {
+	createDerivedAtom,
+	createAtomRegistry,
+	createStoreAtomSelector,
+} from '@wordpress/stan';
 
 /**
  * @typedef {import("../types").WPDataAtomicStoreConfig} WPDataAtomicStoreConfig
@@ -33,28 +37,51 @@ export default function createAtomicStore( name, config ) {
 	return {
 		name,
 		instantiate: ( registry ) => {
-			const selectors = mapValues( config.selectors, ( atomSelector ) => {
-				return ( /** @type {any[]} **/ ...args ) => {
-					const get = registry.__internalGetAtomResolver()
-						? registry.__internalGetAtomResolver()
-						: (
-								/** @type {WPAtom<any>|WPAtomSelector<any>} **/ atom
-						  ) => registry.__internalGetAtomRegistry().get( atom );
-					return get( atomSelector( ...args ) );
-				};
-			} );
+			// Having a dedicated atom registry per store allow us to support
+			// registry inheritance as we won't be retrieving atoms from the wrong registries.
+			const atomRegistry = createAtomRegistry();
+
+			// These are used inside of useSelect when mapSelect can merge selectors
+			// from different stores and different data registries.
+			const registryAtomSelectors = mapValues(
+				config.selectors,
+				( atomSelector ) => {
+					return createStoreAtomSelector(
+						( ...args ) => ( listener ) =>
+							atomRegistry.subscribe(
+								atomSelector( ...args ),
+								listener
+							),
+						( ...args ) => () =>
+							atomRegistry.get( atomSelector( ...args ) ),
+						( ...args ) => ( value ) =>
+							atomRegistry.set( atomSelector( ...args ), value )
+					);
+				}
+			);
+
+			const selectors = mapValues(
+				registryAtomSelectors,
+				( atomSelector, selectorName ) => {
+					return ( /** @type {any[]} **/ ...args ) => {
+						return registry.__internalGetAtomResolver()
+							? registry.__internalGetAtomResolver()(
+									atomSelector( ...args )
+							  )
+							: atomRegistry.get(
+									// @ts-ignore
+									config.selectors[ selectorName ]( ...args )
+							  );
+					};
+				}
+			);
 
 			const actions = mapValues( config.actions, ( atomAction ) => {
 				return ( /** @type {any[]} **/ ...args ) => {
 					return atomAction( ...args )( {
-						get: ( atomCreator ) =>
-							registry
-								.__internalGetAtomRegistry()
-								.get( atomCreator ),
+						get: ( atomCreator ) => atomRegistry.get( atomCreator ),
 						set: ( atomCreator, value ) =>
-							registry
-								.__internalGetAtomRegistry()
-								.set( atomCreator, value ),
+							atomRegistry.set( atomCreator, value ),
 					} );
 				};
 			} );
@@ -73,9 +100,7 @@ export default function createAtomicStore( name, config ) {
 						);
 					} );
 
-					return registry
-						.__internalGetAtomRegistry()
-						.subscribe( atom, listener );
+					return atomRegistry.subscribe( atom, listener );
 				},
 			};
 		},

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

@@ -9,7 +9,7 @@ import {
 	useCallback,
 } from '@wordpress/element';
 import isShallowEqual from '@wordpress/is-shallow-equal';
-import { createDerivedAtom } from '@wordpress/stan';
+import { createDerivedAtom, createAtomRegistry } from '@wordpress/stan';
 import { usePrevious } from '@wordpress/compose';
 
 /**
@@ -60,6 +60,7 @@ import useRegistry from '../registry-provider/use-registry';
  * @return {Function}  A custom react hook.
  */
 export default function useSelect( _mapSelect, deps ) {
+	const atomRegistry = useMemo( () => createAtomRegistry(), [] );
 	const mapSelect = useCallback( _mapSelect, deps );
 	const previousMapSelect = usePrevious( mapSelect );
 	const result = useRef();
@@ -89,7 +90,7 @@ export default function useSelect( _mapSelect, deps ) {
 			},
 			() => {},
 			{ isAsync }
-		)( registry.__internalGetAtomRegistry() );
+		)( atomRegistry );
 	}, [ isAsync, registry, mapSelect ] );
 
 	try {

packages/data/src/redux-store/test/index.js

@@ -1,7 +1,7 @@
 /**
  * WordPress dependencies
  */
-import { createDerivedAtom } from '@wordpress/stan';
+import { createAtomRegistry, createDerivedAtom } from '@wordpress/stan';
 
 /**
  * Internal dependencies
@@ -278,7 +278,7 @@ describe( 'controls', () => {
 		} );
 
 		it( 'should subscribe to atom selectors', async () => {
-			const atomRegistry = registry.__internalGetAtomRegistry();
+			const atomRegistry = createAtomRegistry();
 			const atom = createUseSelectAtom( ( select ) => {
 				return {
 					value: select( 'store1' ).getValue(),
@@ -317,8 +317,7 @@ describe( 'controls', () => {
 					value: select( 'store1' ).getValue(),
 				};
 			} );
-			const atomRegistry = registry.__internalGetAtomRegistry();
-
+			const atomRegistry = createAtomRegistry();
 			const update = jest.fn();
 			const unsubscribe = atomRegistry.subscribe( atom, update );
 			await flushImmediatesAndTicks( 2 );
@@ -344,7 +343,7 @@ describe( 'controls', () => {
 				},
 			} );
 
-			const atomRegistry = registry.__internalGetAtomRegistry();
+			const atomRegistry = createAtomRegistry();
 			const atom = createUseSelectAtom( ( select ) => {
 				return {
 					value: select( 'store2' ).getSubStoreValue(),
@@ -387,7 +386,7 @@ describe( 'controls', () => {
 				},
 			} );
 
-			const atomRegistry = registry.__internalGetAtomRegistry();
+			const atomRegistry = createAtomRegistry();
 			const atom = createUseSelectAtom( ( select ) => {
 				return {
 					value: select( 'store3' ).getSubStoreValue(),

packages/data/src/registry.js

@@ -7,7 +7,7 @@ import memize from 'memize';
 /**
  * WordPress dependencies
  */
-import { createAtomRegistry, createStoreAtom } from '@wordpress/stan';
+import { createStoreAtom } from '@wordpress/stan';
 
 /**
  * Internal dependencies
@@ -55,7 +55,6 @@ import createCoreDataStore from './store';
 export function createRegistry( storeConfigs = {}, parent = null ) {
 	const stores = {};
 	const storesAtoms = {};
-	const atomRegistry = createAtomRegistry();
 	let listeners = [];
 
 	/**
@@ -101,7 +100,7 @@ export function createRegistry( storeConfigs = {}, parent = null ) {
 
 		const store = stores[ storeName ];
 		if ( store ) {
-			// If it's not an atomic store subscribe to the global store.
+			// If it's not an atomic store subscribe to the store.
 			if (
 				! store.__internalIsAtomic &&
 				registry.__internalGetAtomResolver()
@@ -115,11 +114,7 @@ export function createRegistry( storeConfigs = {}, parent = null ) {
 		}
 
 		if ( parent ) {
-			parent.__internalSetAtomResolver(
-				registry.__internalGetAtomResolver()
-			);
-			const ret = parent.select( storeName );
-			return ret;
+			return parent.select( storeName );
 		}
 	}
 
@@ -273,11 +268,11 @@ export function createRegistry( storeConfigs = {}, parent = null ) {
 			return currentAtomResolver;
 		},
 		__internalSetAtomResolver( resolver ) {
+			if ( parent ) {
+				parent.__internalSetAtomResolver( resolver );
+			}
 			currentAtomResolver = resolver;
 		},
-		__internalGetAtomRegistry() {
-			return atomRegistry;
-		},
 	};
 
 	/**

packages/data/src/types.ts

@@ -5,7 +5,6 @@ import type {
 	WPAtom,
 	WPAtomSelector,
 	WPAtomResolver,
-	WPAtomRegistry,
 	WPAtomUpdater,
 } from '@wordpress/stan';
 
@@ -57,11 +56,6 @@ export interface WPDataRegistry {
 	 */
 	register: ( store: WPDataStore ) => void;
 
-	/**
-	 * Retrieves the atom registry.
-	 */
-	__internalGetAtomRegistry: () => WPAtomRegistry;
-
 	/**
 	 * For registry selectors we need to be able to inject the atom resolver.
 	 * This setter/getter allows us to so.

packages/stan/README.md

@@ -197,6 +197,19 @@ _Returns_
 
 -   `WPAtom<T>`: Store Atom.
 
+<a name="createStoreAtomSelector" href="#createStoreAtomSelector">#</a> **createStoreAtomSelector**
+
+_Parameters_
+
+-   _subscribe_ (unknown type): Subscribe to state changes.
+-   _get_ (unknown type): Get the state value.
+-   _dispatch_ (unknown type): Dispatch store changes,
+-   _atomConfig_ `[WPCommonAtomConfig]`: Common Atom config.
+
+_Returns_
+
+-   (unknown type): Atom selector creator.
+
 
 <!-- END TOKEN(Autogenerated API docs) -->
 

packages/stan/src/index.js

@@ -1,7 +1,7 @@
 export { createAtom } from './atom';
 export { createDerivedAtom } from './derived';
-export { createStoreAtom } from './store';
 export { createAtomSelector } from './selector';
+export { createStoreAtom, createStoreAtomSelector } from './store';
 export { createAtomRegistry } from './registry';
 
 /**

packages/stan/src/store.js

@@ -5,6 +5,10 @@
 /**
  * @typedef {import("./types").WPCommonAtomConfig} WPCommonAtomConfig
  */
+/**
+ * @template T
+ * @typedef {import("./types").WPAtomSelector<T>} WPAtomSelector
+ */
 
 /**
  * Creates a store atom.
@@ -37,3 +41,43 @@ export const createStoreAtom = (
 		isResolved: true,
 	};
 };
+
+/**
+ * @template T
+ * @param {(...args:any[]) => (listener: () => void) => (() => void)} subscribe Subscribe to state changes.
+ * @param {(...args:any[]) => () => T}                                get       Get the state value.
+ * @param {(...args:any[]) => (action: any) => void}                  dispatch  Dispatch store changes,
+ * @param {WPCommonAtomConfig=}                                       atomConfig   Common Atom config.
+ * @return {(...args:any[]) => WPAtomSelector<T>} Atom selector creator.
+ */
+export const createStoreAtomSelector = (
+	subscribe,
+	get,
+	dispatch,
+	atomConfig
+) => {
+	const config = {
+		/**
+		 * @param {...any[]} args Selector arguments
+		 * @return {WPAtom<any>} Atom.
+		 */
+		createAtom( ...args ) {
+			return createStoreAtom(
+				subscribe( ...args ),
+				get( ...args ),
+				dispatch( ...args ),
+				{
+					...atomConfig,
+				}
+			);
+		},
+	};
+
+	return ( ...args ) => {
+		return {
+			type: 'selector',
+			config,
+			args,
+		};
+	};
+};