From 936e69fed3d416cf0ccee7b193b4b8a8af130bde Mon Sep 17 00:00:00 2001
From: Daniel Richards <daniel.richards@automattic.com>
Date: Tue, 6 Jul 2021 14:52:37 +0800
Subject: [PATCH] Fix move to widget area checkmark (#33213)

* Traverse through block parents to find widget area

* Refactor to selector

* Switch back to previous selector const

* Add doc block
---
 .../src/filters/move-to-widget-area.js        | 16 +++++++------
 .../hooks/use-last-selected-widget-area.js    | 24 +++++++++----------
 packages/edit-widgets/src/store/selectors.js  | 21 ++++++++++++++++
 3 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/packages/edit-widgets/src/filters/move-to-widget-area.js b/packages/edit-widgets/src/filters/move-to-widget-area.js
index ca7a31cfd493b4..2354c74e52f003 100644
--- a/packages/edit-widgets/src/filters/move-to-widget-area.js
+++ b/packages/edit-widgets/src/filters/move-to-widget-area.js
@@ -6,7 +6,7 @@ import { BlockControls } from '@wordpress/block-editor';
 import { createHigherOrderComponent } from '@wordpress/compose';
 import { useDispatch, useSelect } from '@wordpress/data';
 import { addFilter } from '@wordpress/hooks';
-import { getWidgetIdFromBlock, MoveToWidgetArea } from '@wordpress/widgets';
+import { MoveToWidgetArea } from '@wordpress/widgets';
 
 /**
  * Internal dependencies
@@ -15,8 +15,7 @@ import { store as editWidgetsStore } from '../store';
 
 const withMoveToWidgetAreaToolbarItem = createHigherOrderComponent(
 	( BlockEdit ) => ( props ) => {
-		const widgetId = getWidgetIdFromBlock( props );
-		const blockName = props.name;
+		const { clientId, name: blockName } = props;
 		const {
 			widgetAreas,
 			currentWidgetAreaId,
@@ -29,17 +28,20 @@ const withMoveToWidgetAreaToolbarItem = createHigherOrderComponent(
 				}
 
 				const selectors = select( editWidgetsStore );
+
+				const widgetAreaBlock = selectors.getParentWidgetAreaBlock(
+					clientId
+				);
+
 				return {
 					widgetAreas: selectors.getWidgetAreas(),
-					currentWidgetAreaId: widgetId
-						? selectors.getWidgetAreaForWidgetId( widgetId )?.id
-						: undefined,
+					currentWidgetAreaId: widgetAreaBlock?.attributes?.id,
 					canInsertBlockInWidgetArea: selectors.canInsertBlockInWidgetArea(
 						blockName
 					),
 				};
 			},
-			[ widgetId, blockName ]
+			[ clientId, blockName ]
 		);
 
 		const { moveBlockToWidgetArea } = useDispatch( editWidgetsStore );
diff --git a/packages/edit-widgets/src/hooks/use-last-selected-widget-area.js b/packages/edit-widgets/src/hooks/use-last-selected-widget-area.js
index 4a33e2b3d6097f..ac3e3f3269dc36 100644
--- a/packages/edit-widgets/src/hooks/use-last-selected-widget-area.js
+++ b/packages/edit-widgets/src/hooks/use-last-selected-widget-area.js
@@ -6,6 +6,7 @@ import { useSelect } from '@wordpress/data';
 /**
  * Internal dependencies
  */
+import { store as widgetsEditorStore } from '../store';
 import { buildWidgetAreasPostId, KIND, POST_TYPE } from '../store/utils';
 
 /**
@@ -16,27 +17,24 @@ import { buildWidgetAreasPostId, KIND, POST_TYPE } from '../store/utils';
  */
 const useLastSelectedWidgetArea = () =>
 	useSelect( ( select ) => {
-		const { getBlockSelectionEnd, getBlockParents, getBlockName } = select(
+		const { getBlockSelectionEnd, getBlockName } = select(
 			'core/block-editor'
 		);
-		const blockSelectionEndClientId = getBlockSelectionEnd();
+		const selectionEndClientId = getBlockSelectionEnd();
 
 		// If the selected block is a widget area, return its clientId.
-		if (
-			getBlockName( blockSelectionEndClientId ) === 'core/widget-area'
-		) {
-			return blockSelectionEndClientId;
+		if ( getBlockName( selectionEndClientId ) === 'core/widget-area' ) {
+			return selectionEndClientId;
 		}
 
-		// Otherwise, find the clientId of the top-level widget area by looking
-		// through the selected block's parents.
-		const blockParents = getBlockParents( blockSelectionEndClientId );
-		const rootWidgetAreaClientId = blockParents.find(
-			( clientId ) => getBlockName( clientId ) === 'core/widget-area'
+		const { getParentWidgetAreaBlock } = select( widgetsEditorStore );
+		const widgetAreaBlock = getParentWidgetAreaBlock(
+			selectionEndClientId
 		);
+		const widgetAreaBlockClientId = widgetAreaBlock?.clientId;
 
-		if ( rootWidgetAreaClientId ) {
-			return rootWidgetAreaClientId;
+		if ( widgetAreaBlockClientId ) {
+			return widgetAreaBlockClientId;
 		}
 
 		// If no widget area has been selected, return the clientId of the first
diff --git a/packages/edit-widgets/src/store/selectors.js b/packages/edit-widgets/src/store/selectors.js
index 7e30ed3f81f0cf..8dace5020abf52 100644
--- a/packages/edit-widgets/src/store/selectors.js
+++ b/packages/edit-widgets/src/store/selectors.js
@@ -77,6 +77,27 @@ export const getWidgetAreaForWidgetId = createRegistrySelector(
 	}
 );
 
+/**
+ * Given a child client id, returns the parent widget area block.
+ *
+ * @param {string} clientId The client id of a block in a widget area.
+ *
+ * @return {WPBlock} The widget area block.
+ */
+export const getParentWidgetAreaBlock = createRegistrySelector(
+	( select ) => ( state, clientId ) => {
+		const { getBlock, getBlockName, getBlockParents } = select(
+			blockEditorStore
+		);
+		const blockParents = getBlockParents( clientId );
+		const widgetAreaClientId = blockParents.find(
+			( parentClientId ) =>
+				getBlockName( parentClientId ) === 'core/widget-area'
+		);
+		return getBlock( widgetAreaClientId );
+	}
+);
+
 export const getEditedWidgetAreas = createRegistrySelector(
 	( select ) => ( state, ids ) => {
 		let widgetAreas = select( editWidgetsStoreName ).getWidgetAreas();