New useBlockElementRef hook for storing block element into a ref

packages/block-editor/src/components/block-breadcrumb/index.js

@@ -5,14 +5,15 @@ import { Button } from '@wordpress/components';
 import { useSelect, useDispatch } from '@wordpress/data';
 import { __ } from '@wordpress/i18n';
 import { chevronRightSmall, Icon } from '@wordpress/icons';
+import { useRef } from '@wordpress/element';
 
 /**
  * Internal dependencies
  */
 import BlockTitle from '../block-title';
 import { store as blockEditorStore } from '../../store';
 import { unlock } from '../../lock-unlock';
-import { __unstableUseBlockRef as useBlockRef } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElementRef } from '../block-list/use-block-props/use-block-refs';
 import getEditorRegion from '../../utils/get-editor-region';
 
 /**
@@ -41,7 +42,8 @@ function BlockBreadcrumb( { rootLabelText } ) {
 
 	// We don't care about this specific ref, but this is a way
 	// to get a ref within the editor canvas so we can focus it later.
-	const blockRef = useBlockRef( clientId );
+	const blockRef = useRef();
+	useBlockElementRef( clientId, blockRef );
 
 	/*
 	 * Disable reason: The `list` ARIA role is redundant but

packages/block-editor/src/components/block-draggable/index.js

@@ -13,7 +13,7 @@ import { throttle } from '@wordpress/compose';
 import BlockDraggableChip from './draggable-chip';
 import useScrollWhenDragging from './use-scroll-when-dragging';
 import { store as blockEditorStore } from '../../store';
-import { __unstableUseBlockRef as useBlockRef } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import { isDropTargetValid } from '../use-block-drop-zone';
 
 const BlockDraggable = ( {
@@ -82,8 +82,8 @@ const BlockDraggable = ( {
 	}, [] );
 
 	// Find the root of the editor iframe.
-	const blockRef = useBlockRef( clientIds[ 0 ] );
-	const editorRoot = blockRef.current?.closest( 'body' );
+	const blockEl = useBlockElement( clientIds[ 0 ] );
+	const editorRoot = blockEl?.closest( 'body' );
 
 	/*
 	 * Add a dragover event listener to the editor root to track the blocks being dragged over.

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

@@ -3,10 +3,9 @@
  */
 import {
 	useContext,
-	useMemo,
-	useRef,
 	useState,
 	useLayoutEffect,
+	useCallback,
 } from '@wordpress/element';
 import { useRefEffect } from '@wordpress/compose';
 
@@ -16,7 +15,7 @@ import { useRefEffect } from '@wordpress/compose';
 import { BlockRefs } from '../../provider/block-refs-provider';
 
 /** @typedef {import('@wordpress/element').RefCallback} RefCallback */
-/** @typedef {import('@wordpress/element').RefObject} RefObject */
+/** @typedef {import('@wordpress/element').Ref} Ref */
 
 /**
  * Provides a ref to the BlockRefs context.
@@ -37,30 +36,35 @@ export function useBlockRefProvider( clientId ) {
 }
 
 /**
- * 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.
+ * Tracks the DOM element for the block identified by `clientId` and assigns it to the `ref`
+ * whenever it changes.
  *
- * @param {string} clientId The client ID to get a ref for.
- *
- * @return {RefObject} A ref containing the element.
+ * @param {string} clientId The client ID to track.
+ * @param {Ref}    ref      The ref object/callback to assign to.
  */
-function useBlockRef( clientId ) {
+export function useBlockElementRef( clientId, ref ) {
 	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() {
-				return refsMap.get( latestClientId.current ) ?? null;
-			},
-		} ),
-		[ refsMap ]
+	const setRef = useCallback(
+		( el ) => {
+			if ( typeof ref === 'function' ) {
+				ref( el );
+			} else if ( ref ) {
+				ref.current = el;
+			}
+		},
+		[ ref ]
 	);
+
+	useLayoutEffect( () => {
+		setRef( refsMap.get( clientId ) );
+		const unsubscribe = refsMap.subscribe( clientId, () =>
+			setRef( refsMap.get( clientId ) )
+		);
+		return () => {
+			unsubscribe();
+			setRef( null );
+		};
+	}, [ refsMap, clientId, setRef ] );
 }
 
 /**
@@ -71,20 +75,8 @@ function useBlockRef( clientId ) {
  *
  * @return {Element|null} The block's wrapper element.
  */
-function useBlockElement( clientId ) {
-	const { refsMap } = useContext( BlockRefs );
+export function useBlockElement( clientId ) {
 	const [ blockElement, setBlockElement ] = useState( null );
-	// Delay setting the resulting `blockElement` until an effect. If the block element
-	// changes (i.e., the block is unmounted and re-mounted), this allows enough time
-	// for the ref callbacks to clean up the old element and set the new one.
-	useLayoutEffect( () => {
-		setBlockElement( refsMap.get( clientId ) );
-		return refsMap.subscribe( clientId, () =>
-			setBlockElement( refsMap.get( clientId ) )
-		);
-	}, [ refsMap, clientId ] );
+	useBlockElementRef( clientId, setBlockElement );
 	return blockElement;
 }
-
-export { useBlockRef as __unstableUseBlockRef };
-export { useBlockElement as __unstableUseBlockElement };

packages/block-editor/src/components/block-popover/cover.js

@@ -6,7 +6,7 @@ import { useEffect, useState, useMemo, forwardRef } from '@wordpress/element';
 /**
  * Internal dependencies
  */
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import { PrivateBlockPopover } from '.';
 
 function BlockPopoverCover(

packages/block-editor/src/components/block-popover/inbetween.js

@@ -20,7 +20,7 @@ import { isRTL } from '@wordpress/i18n';
  * Internal dependencies
  */
 import { store as blockEditorStore } from '../../store';
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import usePopoverScroll from './use-popover-scroll';
 
 const MAX_POPOVER_RECOMPUTE_COUNTER = Number.MAX_SAFE_INTEGER;

packages/block-editor/src/components/block-popover/index.js

@@ -18,7 +18,7 @@ import {
 /**
  * Internal dependencies
  */
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import usePopoverScroll from './use-popover-scroll';
 
 const MAX_POPOVER_RECOMPUTE_COUNTER = Number.MAX_SAFE_INTEGER;

packages/block-editor/src/components/block-tools/block-selection-button.js

@@ -37,7 +37,7 @@ import BlockTitle from '../block-title';
 import BlockIcon from '../block-icon';
 import { store as blockEditorStore } from '../../store';
 import BlockDraggable from '../block-draggable';
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 
 /**
  * Block selection button component, displaying the label of the block. If the block

packages/block-editor/src/components/block-tools/use-block-toolbar-popover-props.js

@@ -15,7 +15,7 @@ import {
  * Internal dependencies
  */
 import { store as blockEditorStore } from '../../store';
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import { hasStickyOrFixedPositionValue } from '../../hooks/position';
 
 const COMMON_PROPS = {

packages/block-editor/src/components/grid/grid-item-resizer.js

@@ -7,7 +7,7 @@ import { useState, useEffect } from '@wordpress/element';
 /**
  * Internal dependencies
  */
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import BlockPopoverCover from '../block-popover/cover';
 import { getComputedCSS, getGridTracks, getClosestTrack } from './utils';
 

packages/block-editor/src/components/grid/grid-visualizer.js

@@ -13,7 +13,7 @@ import { __experimentalUseDropZone as useDropZone } from '@wordpress/compose';
 /**
  * Internal dependencies
  */
-import { __unstableUseBlockElement as useBlockElement } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../block-list/use-block-props/use-block-refs';
 import BlockPopoverCover from '../block-popover/cover';
 import { range, GridRect, getGridInfo } from './utils';
 import { store as blockEditorStore } from '../../store';

packages/block-editor/src/components/skip-to-selected-block/index.js

@@ -4,12 +4,13 @@
 import { useSelect } from '@wordpress/data';
 import { __ } from '@wordpress/i18n';
 import { Button } from '@wordpress/components';
+import { useRef } from '@wordpress/element';
 
 /**
  * Internal dependencies
  */
 import { store as blockEditorStore } from '../../store';
-import { __unstableUseBlockRef as useBlockRef } from '../block-list/use-block-props/use-block-refs';
+import { useBlockElementRef } from '../block-list/use-block-props/use-block-refs';
 
 /**
  * @see https://github.com/WordPress/gutenberg/blob/HEAD/packages/block-editor/src/components/skip-to-selected-block/README.md
@@ -19,9 +20,10 @@ export default function SkipToSelectedBlock() {
 		( select ) => select( blockEditorStore ).getBlockSelectionStart(),
 		[]
 	);
-	const ref = useBlockRef( selectedBlockClientId );
+	const ref = useRef();
+	useBlockElementRef( selectedBlockClientId, ref );
 	const onClick = () => {
-		ref.current.focus();
+		ref.current?.focus();
 	};
 
 	return selectedBlockClientId ? (

packages/block-editor/src/hooks/contrast-checker.js

@@ -7,7 +7,7 @@ import { useState, useEffect } from '@wordpress/element';
  * Internal dependencies
  */
 import ContrastChecker from '../components/contrast-checker';
-import { __unstableUseBlockRef as useBlockRef } from '../components/block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../components/block-list/use-block-props/use-block-refs';
 
 function getComputedStyle( node ) {
 	return node.ownerDocument.defaultView.getComputedStyle( node );
@@ -17,23 +17,23 @@ export default function BlockColorContrastChecker( { clientId } ) {
 	const [ detectedBackgroundColor, setDetectedBackgroundColor ] = useState();
 	const [ detectedColor, setDetectedColor ] = useState();
 	const [ detectedLinkColor, setDetectedLinkColor ] = useState();
-	const ref = useBlockRef( clientId );
+	const blockEl = useBlockElement( clientId );
 
 	// There are so many things that can change the color of a block
 	// So we perform this check on every render.
 	// eslint-disable-next-line react-hooks/exhaustive-deps
 	useEffect( () => {
-		if ( ! ref.current ) {
+		if ( ! blockEl ) {
 			return;
 		}
-		setDetectedColor( getComputedStyle( ref.current ).color );
+		setDetectedColor( getComputedStyle( blockEl ).color );
 
-		const firstLinkElement = ref.current?.querySelector( 'a' );
+		const firstLinkElement = blockEl.querySelector( 'a' );
 		if ( firstLinkElement && !! firstLinkElement.innerText ) {
 			setDetectedLinkColor( getComputedStyle( firstLinkElement ).color );
 		}
 
-		let backgroundColorNode = ref.current;
+		let backgroundColorNode = blockEl;
 		let backgroundColor =
 			getComputedStyle( backgroundColorNode ).backgroundColor;
 		while (
@@ -48,7 +48,7 @@ export default function BlockColorContrastChecker( { clientId } ) {
 		}
 
 		setDetectedBackgroundColor( backgroundColor );
-	} );
+	}, [ blockEl ] );
 
 	return (
 		<ContrastChecker

packages/block-editor/src/hooks/duotone.js

@@ -35,7 +35,7 @@ import { scopeSelector } from '../components/global-styles/utils';
 import { useBlockSettings, useStyleOverride } from './utils';
 import { default as StylesFiltersPanel } from '../components/global-styles/filters-panel';
 import { useBlockEditingMode } from '../components/block-editing-mode';
-import { __unstableUseBlockElement as useBlockElement } from '../components/block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../components/block-list/use-block-props/use-block-refs';
 
 const EMPTY_ARRAY = [];
 

packages/block-editor/src/hooks/spacing-visualizer.js

@@ -14,7 +14,7 @@ import isShallowEqual from '@wordpress/is-shallow-equal';
  * Internal dependencies
  */
 import BlockPopoverCover from '../components/block-popover/cover';
-import { __unstableUseBlockElement as useBlockElement } from '../components/block-list/use-block-props/use-block-refs';
+import { useBlockElement } from '../components/block-list/use-block-props/use-block-refs';
 
 function SpacingVisualizer( { clientId, value, computeStyle, forceShow } ) {
 	const blockElement = useBlockElement( clientId );