From f7b42299d0f3320bf8339aad410dc75e52b6aa50 Mon Sep 17 00:00:00 2001 From: Jarda Snajdr Date: Tue, 16 Jul 2024 18:17:43 +0200 Subject: [PATCH] useBlockElement: return null until ref callback has time to clean up the old element (#63565) Unlinked contributors: sergiu-radu. Co-authored-by: jsnajdr Co-authored-by: talldan Co-authored-by: ciampo Co-authored-by: ellatrix --- .../use-block-props/use-block-refs.js | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js b/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js index 297aed456df7fc..056ade045d1654 100644 --- a/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js +++ b/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 };