From f48d2fda4cdc694a15bdf51401125a3020ccb09f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ella=20van=C2=A0Durpe?= <ella@vandurpe.com>
Date: Mon, 17 May 2021 21:17:44 +0300
Subject: [PATCH 1/3] Fix block ref when multiple useBlockProps hooks are
 called

---
 .../use-block-props/use-block-refs.js         | 41 ++++++++++++-------
 1 file changed, 27 insertions(+), 14 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 50794c611f1f68..a25fec62a92923 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
@@ -29,21 +29,24 @@ export function useBlockRefProvider( clientId ) {
 	const { refs, callbacks } = useContext( BlockRefs );
 	const ref = useRef();
 	useLayoutEffect( () => {
-		refs.set( clientId, ref );
+		refs.set( ref, clientId );
 		return () => {
-			refs.delete( clientId );
+			refs.delete( ref );
 		};
-	}, [] );
-	return useCallback( ( element ) => {
-		// Update the ref in the provider.
-		ref.current = element;
-		// Call any update functions.
-		callbacks.forEach( ( id, setElement ) => {
-			if ( clientId === id ) {
-				setElement( element );
-			}
-		} );
-	}, [] );
+	}, [ clientId ] );
+	return useCallback(
+		( element ) => {
+			// Update the ref in the provider.
+			ref.current = element;
+			// Call any update functions.
+			callbacks.forEach( ( id, setElement ) => {
+				if ( clientId === id ) {
+					setElement( element );
+				}
+			} );
+		},
+		[ clientId ]
+	);
 }
 
 /**
@@ -63,7 +66,17 @@ function useBlockRef( clientId ) {
 	return useMemo(
 		() => ( {
 			get current() {
-				return refs.get( freshClientId.current )?.current || null;
+				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;
 			},
 		} ),
 		[]

From 1ea9cd6738947c6644c8a909ed4954e5a7de6e7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ella=20van=C2=A0Durpe?= <ella@vandurpe.com>
Date: Wed, 19 May 2021 17:25:11 +0300
Subject: [PATCH 2/3] Don't call callbacks with null

---
 .../block-list/use-block-props/use-block-refs.js         | 9 ++++-----
 1 file changed, 4 insertions(+), 5 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 a25fec62a92923..44da41fe8bf91e 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
@@ -2,13 +2,13 @@
  * WordPress dependencies
  */
 import {
-	useCallback,
 	useContext,
 	useLayoutEffect,
 	useMemo,
 	useRef,
 	useState,
 } from '@wordpress/element';
+import { useRefEffect } from '@wordpress/compose';
 
 /**
  * Internal dependencies
@@ -34,7 +34,7 @@ export function useBlockRefProvider( clientId ) {
 			refs.delete( ref );
 		};
 	}, [ clientId ] );
-	return useCallback(
+	return useRefEffect(
 		( element ) => {
 			// Update the ref in the provider.
 			ref.current = element;
@@ -93,8 +93,7 @@ function useBlockRef( clientId ) {
  */
 function useBlockElement( clientId ) {
 	const { callbacks } = useContext( BlockRefs );
-	const ref = useBlockRef( clientId );
-	const [ element, setElement ] = useState( null );
+	const [ element, setElement ] = useState( useBlockRef( clientId ).current );
 
 	useLayoutEffect( () => {
 		if ( ! clientId ) {
@@ -107,7 +106,7 @@ function useBlockElement( clientId ) {
 		};
 	}, [ clientId ] );
 
-	return ref.current || element;
+	return element;
 }
 
 export { useBlockRef as __unstableUseBlockRef };

From a07718bab720e2d0ee6f44d891c3841cf1485d26 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ella=20van=C2=A0Durpe?= <ella@vandurpe.com>
Date: Wed, 19 May 2021 19:21:45 +0300
Subject: [PATCH 3/3] Revert changes to useBlockElement

---
 .../components/block-list/use-block-props/use-block-refs.js  | 5 +++--
 1 file changed, 3 insertions(+), 2 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 44da41fe8bf91e..56e424319739e1 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
@@ -93,7 +93,8 @@ function useBlockRef( clientId ) {
  */
 function useBlockElement( clientId ) {
 	const { callbacks } = useContext( BlockRefs );
-	const [ element, setElement ] = useState( useBlockRef( clientId ).current );
+	const ref = useBlockRef( clientId );
+	const [ element, setElement ] = useState( null );
 
 	useLayoutEffect( () => {
 		if ( ! clientId ) {
@@ -106,7 +107,7 @@ function useBlockElement( clientId ) {
 		};
 	}, [ clientId ] );
 
-	return element;
+	return ref.current || element;
 }
 
 export { useBlockRef as __unstableUseBlockRef };