useBlockElement: return null until ref callback has time to clean up …

…the old element (#63565)

Unlinked contributors: sergiu-radu.

Co-authored-by: jsnajdr <[email protected]>
Co-authored-by: talldan <[email protected]>
Co-authored-by: ciampo <[email protected]>
Co-authored-by: ellatrix <[email protected]>

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

@@ -1,8 +1,14 @@
 /**
  * WordPress dependencies
  */
-import { useContext, useMemo, useRef } from '@wordpress/element';
-import { useRefEffect, useObservableValue } from '@wordpress/compose';
+import {
+	useContext,
+	useMemo,
+	useRef,
+	useState,
+	useLayoutEffect,
+} from '@wordpress/element';
+import { useRefEffect } from '@wordpress/compose';
 
 /**
  * Internal dependencies
@@ -67,7 +73,17 @@ function useBlockRef( clientId ) {
  */
 function useBlockElement( clientId ) {
 	const { refsMap } = useContext( BlockRefs );
-	return useObservableValue( refsMap, clientId ) ?? null;
+	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 ] );
+	return blockElement;
 }
 
 export { useBlockRef as __unstableUseBlockRef };