useBlockRefs: use more efficient lookup map, use uSES (#60945)

* useBlockRefs: use more efficient lookup map, use uSES

* Rewrite block refs with observableMap, which moves to compose

* Improve docs

* Add changelog entry

Author: jsnajdr

packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js

@@ -1,14 +1,8 @@
 /**
  * WordPress dependencies
  */
-import {
-	useContext,
-	useLayoutEffect,
-	useMemo,
-	useRef,
-	useState,
-} from '@wordpress/element';
-import { useRefEffect } from '@wordpress/compose';
+import { useContext, useMemo, useRef } from '@wordpress/element';
+import { useRefEffect, useObservableValue } from '@wordpress/compose';
 
 /**
  * Internal dependencies
@@ -26,60 +20,40 @@ import { BlockRefs } from '../../provider/block-refs-provider';
  * @return {RefCallback} Ref callback.
  */
 export function useBlockRefProvider( clientId ) {
-	const { refs, callbacks } = useContext( BlockRefs );
-	const ref = useRef();
-	useLayoutEffect( () => {
-		refs.set( ref, clientId );
-		return () => {
-			refs.delete( ref );
-		};
-	}, [ clientId ] );
+	const { refsMap } = useContext( BlockRefs );
 	return useRefEffect(
 		( element ) => {
-			// Update the ref in the provider.
-			ref.current = element;
-			// Call any update functions.
-			callbacks.forEach( ( id, setElement ) => {
-				if ( clientId === id ) {
-					setElement( element );
-				}
-			} );
+			refsMap.set( clientId, element );
+			return () => refsMap.delete( clientId );
 		},
 		[ clientId ]
 	);
 }
 
 /**
- * Gets a ref pointing to the current block element. Continues to return a
- * stable ref even if the block client ID changes.
+ * Gets a ref pointing to the current block element. Continues to return the same
+ * stable ref object even if the `clientId` argument changes. This hook is not
+ * reactive, i.e., it won't trigger a rerender of the calling component if the
+ * ref value changes. For reactive use cases there is the `useBlockElement` hook.
  *
  * @param {string} clientId The client ID to get a ref for.
  *
  * @return {RefObject} A ref containing the element.
  */
 function useBlockRef( clientId ) {
-	const { refs } = useContext( BlockRefs );
-	const freshClientId = useRef();
-	freshClientId.current = clientId;
+	const { refsMap } = useContext( BlockRefs );
+	const latestClientId = useRef();
+	latestClientId.current = clientId;
+
 	// Always return an object, even if no ref exists for a given client ID, so
 	// that `current` works at a later point.
 	return useMemo(
 		() => ( {
 			get current() {
-				let element = null;
-
-				// Multiple refs may be created for a single block. Find the
-				// first that has an element set.
-				for ( const [ ref, id ] of refs.entries() ) {
-					if ( id === freshClientId.current && ref.current ) {
-						element = ref.current;
-					}
-				}
-
-				return element;
+				return refsMap.get( latestClientId.current ) ?? null;
 			},
 		} ),
-		[]
+		[ refsMap ]
 	);
 }
 
@@ -92,22 +66,8 @@ function useBlockRef( clientId ) {
  * @return {Element|null} The block's wrapper element.
  */
 function useBlockElement( clientId ) {
-	const { callbacks } = useContext( BlockRefs );
-	const ref = useBlockRef( clientId );
-	const [ element, setElement ] = useState( null );
-
-	useLayoutEffect( () => {
-		if ( ! clientId ) {
-			return;
-		}
-
-		callbacks.set( setElement, clientId );
-		return () => {
-			callbacks.delete( setElement );
-		};
-	}, [ clientId ] );
-
-	return ref.current || element;
+	const { refsMap } = useContext( BlockRefs );
+	return useObservableValue( refsMap, clientId ) ?? null;
 }
 
 export { useBlockRef as __unstableUseBlockRef };

packages/block-editor/src/components/provider/block-refs-provider.js

@@ -2,17 +2,12 @@
  * WordPress dependencies
  */
 import { createContext, useMemo } from '@wordpress/element';
+import { observableMap } from '@wordpress/compose';
 
-export const BlockRefs = createContext( {
-	refs: new Map(),
-	callbacks: new Map(),
-} );
+export const BlockRefs = createContext( { refsMap: observableMap() } );
 
 export function BlockRefsProvider( { children } ) {
-	const value = useMemo(
-		() => ( { refs: new Map(), callbacks: new Map() } ),
-		[]
-	);
+	const value = useMemo( () => ( { refsMap: observableMap() } ), [] );
 	return (
 		<BlockRefs.Provider value={ value }>{ children }</BlockRefs.Provider>
 	);

packages/components/src/slot-fill/bubbles-virtually/slot-fill-context.ts

@@ -3,11 +3,12 @@
  */
 import { createContext } from '@wordpress/element';
 import warning from '@wordpress/warning';
+import { observableMap } from '@wordpress/compose';
+
 /**
  * Internal dependencies
  */
 import type { SlotFillBubblesVirtuallyContext } from '../types';
-import { observableMap } from './observable-map';
 
 const initialContextValue: SlotFillBubblesVirtuallyContext = {
 	slots: observableMap(),

packages/components/src/slot-fill/bubbles-virtually/slot-fill-provider.tsx

@@ -3,6 +3,7 @@
  */
 import { useMemo } from '@wordpress/element';
 import isShallowEqual from '@wordpress/is-shallow-equal';
+import { observableMap } from '@wordpress/compose';
 
 /**
  * Internal dependencies
@@ -12,7 +13,6 @@ import type {
 	SlotFillProviderProps,
 	SlotFillBubblesVirtuallyContext,
 } from '../types';
-import { observableMap } from './observable-map';
 
 function createSlotRegistry(): SlotFillBubblesVirtuallyContext {
 	const slots: SlotFillBubblesVirtuallyContext[ 'slots' ] = observableMap();

packages/components/src/slot-fill/bubbles-virtually/use-slot-fills.ts

@@ -2,13 +2,13 @@
  * WordPress dependencies
  */
 import { useContext } from '@wordpress/element';
+import { useObservableValue } from '@wordpress/compose';
 
 /**
  * Internal dependencies
  */
 import SlotFillContext from './slot-fill-context';
 import type { SlotKey } from '../types';
-import { useObservableValue } from './observable-map';
 
 export default function useSlotFills( name: SlotKey ) {
 	const registry = useContext( SlotFillContext );

packages/components/src/slot-fill/bubbles-virtually/use-slot.ts

@@ -2,6 +2,7 @@
  * WordPress dependencies
  */
 import { useMemo, useContext } from '@wordpress/element';
+import { useObservableValue } from '@wordpress/compose';
 
 /**
  * Internal dependencies
@@ -13,7 +14,6 @@ import type {
 	FillProps,
 	SlotKey,
 } from '../types';
-import { useObservableValue } from './observable-map';
 
 export default function useSlot( name: SlotKey ) {
 	const registry = useContext( SlotFillContext );

packages/components/src/slot-fill/types.ts

@@ -4,9 +4,9 @@
 import type { Component, MutableRefObject, ReactNode, RefObject } from 'react';
 
 /**
- * Internal dependencies
+ * WordPress dependencies
  */
-import type { ObservableMap } from './bubbles-virtually/observable-map';
+import type { ObservableMap } from '@wordpress/compose';
 
 export type DistributiveOmit< T, K extends keyof any > = T extends any
 	? Omit< T, K >

packages/compose/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+-   Added new `observableMap` data structure and `useObservableValue` React hook ([#60945](https://github.com/WordPress/gutenberg/pull/60945)).
+
 ## 6.33.0 (2024-04-19)
 
 ## 6.32.0 (2024-04-03)

packages/compose/README.md

@@ -129,6 +129,14 @@ _Returns_
 
 -   Higher-order component.
 
+### observableMap
+
+A constructor (factory) for `ObservableMap`, a map-like key/value data structure where the individual entries are observable: using the `subscribe` method, you can subscribe to updates for a particular keys. Each subscriber always observes one specific key and is not notified about any unrelated changes (for different keys) in the `ObservableMap`.
+
+_Returns_
+
+-   `ObservableMap< K, V >`: A new instance of the `ObservableMap` type.
+
 ### pipe
 
 Composes multiple higher-order components into a single higher-order component. Performs left-to-right function composition, where each successive invocation is supplied the return value of the previous.
@@ -442,6 +450,19 @@ _Returns_
 
 -   `import('react').RefCallback<TypeFromRef<TRef>>`: The merged ref callback.
 
+### useObservableValue
+
+React hook that lets you observe an entry in an `ObservableMap`. The hook returns the current value corresponding to the key, or `undefined` when there is no value stored. It also observes changes to the value and triggers an update of the calling component in case the value changes.
+
+_Parameters_
+
+-   _map_ `ObservableMap< K, V >`: The `ObservableMap` to observe.
+-   _name_ `K`: The map key to observe.
+
+_Returns_
+
+-   `V | undefined`: The value corresponding to the map key requested.
+
 ### usePrevious
 
 Use something's value from the previous render. Based on <https://usehooks.com/usePrevious/>.

packages/compose/src/hooks/use-observable-value/index.ts

@@ -0,0 +1,35 @@
+/**
+ * WordPress dependencies
+ */
+import { useMemo, useSyncExternalStore } from '@wordpress/element';
+
+/**
+ * Internal dependencies
+ */
+import type { ObservableMap } from '../../utils/observable-map';
+
+/**
+ * React hook that lets you observe an entry in an `ObservableMap`. The hook returns the
+ * current value corresponding to the key, or `undefined` when there is no value stored.
+ * It also observes changes to the value and triggers an update of the calling component
+ * in case the value changes.
+ *
+ * @template K    The type of the keys in the map.
+ * @template V    The type of the values in the map.
+ * @param    map  The `ObservableMap` to observe.
+ * @param    name The map key to observe.
+ * @return   The value corresponding to the map key requested.
+ */
+export default function useObservableValue< K, V >(
+	map: ObservableMap< K, V >,
+	name: K
+): V | undefined {
+	const [ subscribe, getValue ] = useMemo(
+		() => [
+			( listener: () => void ) => map.subscribe( name, listener ),
+			() => map.get( name ),
+		],
+		[ map, name ]
+	);
+	return useSyncExternalStore( subscribe, getValue, getValue );
+}

packages/compose/src/hooks/use-observable-value/test/index.js

@@ -0,0 +1,42 @@
+/**
+ * External dependencies
+ */
+import { render, screen, act } from '@testing-library/react';
+
+/**
+ * Internal dependencies
+ */
+import { observableMap } from '../../../utils/observable-map';
+import useObservableValue from '..';
+
+describe( 'useObservableValue', () => {
+	test( 'reacts only to the specified key', () => {
+		const map = observableMap();
+		map.set( 'a', 1 );
+
+		const MapUI = jest.fn( () => {
+			const value = useObservableValue( map, 'a' );
+			return <div>value is { value }</div>;
+		} );
+
+		render( <MapUI /> );
+		expect( screen.getByText( /^value is/ ) ).toHaveTextContent(
+			'value is 1'
+		);
+		expect( MapUI ).toHaveBeenCalledTimes( 1 );
+
+		act( () => {
+			map.set( 'a', 2 );
+		} );
+		expect( screen.getByText( /^value is/ ) ).toHaveTextContent(
+			'value is 2'
+		);
+		expect( MapUI ).toHaveBeenCalledTimes( 2 );
+
+		// check that setting unobserved map key doesn't trigger a render at all
+		act( () => {
+			map.set( 'b', 1 );
+		} );
+		expect( MapUI ).toHaveBeenCalledTimes( 2 );
+	} );
+} );

packages/compose/src/index.js

@@ -4,6 +4,8 @@ export * from './utils/create-higher-order-component';
 export * from './utils/debounce';
 // The `throttle` helper and its types.
 export * from './utils/throttle';
+// The `ObservableMap` data structure
+export * from './utils/observable-map';
 
 // The `compose` and `pipe` helpers (inspired by `flowRight` and `flow` from Lodash).
 export { default as compose } from './higher-order/compose';
@@ -46,3 +48,4 @@ export { default as useRefEffect } from './hooks/use-ref-effect';
 export { default as __experimentalUseDropZone } from './hooks/use-drop-zone';
 export { default as useFocusableIframe } from './hooks/use-focusable-iframe';
 export { default as __experimentalUseFixedWindowList } from './hooks/use-fixed-window-list';
+export { default as useObservableValue } from './hooks/use-observable-value';

packages/compose/src/index.native.js

@@ -4,6 +4,8 @@ export * from './utils/create-higher-order-component';
 export * from './utils/debounce';
 // The `throttle` helper and its types.
 export * from './utils/throttle';
+// The `ObservableMap` data structure
+export * from './utils/observable-map';
 
 // The `compose` and `pipe` helpers (inspired by `flowRight` and `flow` from Lodash).
 export { default as compose } from './higher-order/compose';
@@ -39,3 +41,4 @@ export { default as useThrottle } from './hooks/use-throttle';
 export { default as useMergeRefs } from './hooks/use-merge-refs';
 export { default as useRefEffect } from './hooks/use-ref-effect';
 export { default as useNetworkConnectivity } from './hooks/use-network-connectivity';
+export { default as useObservableValue } from './hooks/use-observable-value';