Fix controlled inner blocks parent block attributes updates (#35827)

* Save controlled inner blocks wrapper attributes

* Try simplifying the logic

* Clarify function calls

* Ensure `UPDATE_BLOCK` reducer maintains tree structure for updated block

Co-authored-by: Riad Benguella <benguella@gmail.com>

packages/block-editor/src/store/reducer.js

@@ -240,32 +240,40 @@ function buildBlockTree( state, blocks ) {
 	return result;
 }
 
-function updateParentInnerBlocksInTree( state, tree, updatedClientIds ) {
-	const clientIds = new Set( [] );
+function updateParentInnerBlocksInTree(
+	state,
+	tree,
+	updatedClientIds,
+	updateChildrenOfUpdatedClientIds = false
+) {
+	const uncontrolledParents = new Set( [] );
 	const controlledParents = new Set();
 	for ( const clientId of updatedClientIds ) {
-		let current = clientId;
+		let current = updateChildrenOfUpdatedClientIds
+			? clientId
+			: state.parents[ clientId ];
 		do {
 			if ( state.controlledInnerBlocks[ current ] ) {
-				controlledParents.add( current );
+				// Should stop on controlled blocks.
 				// If we reach a controlled parent, break out of the loop.
+				controlledParents.add( current );
 				break;
 			} else {
-				clientIds.add( current );
+				// else continue traversing up through parents.
+				uncontrolledParents.add( current );
+				current = state.parents[ current ];
 			}
-			// Should stop on controlled blocks.
-			current = state.parents[ current ];
 		} while ( current !== undefined );
 	}
 
 	// To make sure the order of assignments doesn't matter,
 	// we first create empty objects and mutates the inner blocks later.
-	for ( const clientId of clientIds ) {
+	for ( const clientId of uncontrolledParents ) {
 		tree[ clientId ] = {
 			...tree[ clientId ],
 		};
 	}
-	for ( const clientId of clientIds ) {
+	for ( const clientId of uncontrolledParents ) {
 		tree[ clientId ].innerBlocks = ( state.order[ clientId ] || [] ).map(
 			( subClientId ) => tree[ subClientId ]
 		);
@@ -311,7 +319,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 					...newState.tree,
 					...subTree,
 				},
-				action.rootClientId ? [ action.rootClientId ] : [ '' ]
+				action.rootClientId ? [ action.rootClientId ] : [ '' ],
+				true
 			);
 			break;
 		}
@@ -321,11 +330,13 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 				{
 					...newState.tree,
 					[ action.clientId ]: {
+						...newState.tree[ action.clientId ],
 						...newState.byClientId[ action.clientId ],
 						attributes: newState.attributes[ action.clientId ],
 					},
 				},
-				[ action.clientId ]
+				[ action.clientId ],
+				false
 			);
 			break;
 		case 'UPDATE_BLOCK_ATTRIBUTES': {
@@ -345,7 +356,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 					...newState.tree,
 					...newSubTree,
 				},
-				action.clientIds
+				action.clientIds,
+				false
 			);
 			break;
 		}
@@ -364,7 +376,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 					),
 					...subTree,
 				},
-				action.blocks.map( ( b ) => b.clientId )
+				action.blocks.map( ( b ) => b.clientId ),
+				false
 			);
 
 			// If there are no replaced blocks, it means we're removing blocks so we need to update their parent.
@@ -381,7 +394,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 			newState.tree = updateParentInnerBlocksInTree(
 				newState,
 				newState.tree,
-				parentsOfRemovedBlocks
+				parentsOfRemovedBlocks,
+				true
 			);
 			break;
 		}
@@ -406,7 +420,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 						)
 					)
 				),
-				parentsOfRemovedBlocks
+				parentsOfRemovedBlocks,
+				true
 			);
 			break;
 		case 'MOVE_BLOCKS_TO_POSITION': {
@@ -423,7 +438,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 			newState.tree = updateParentInnerBlocksInTree(
 				newState,
 				newState.tree,
-				updatedBlockUids
+				updatedBlockUids,
+				true
 			);
 			break;
 		}
@@ -435,7 +451,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 			newState.tree = updateParentInnerBlocksInTree(
 				newState,
 				newState.tree,
-				updatedBlockUids
+				updatedBlockUids,
+				true
 			);
 			break;
 		}
@@ -462,7 +479,8 @@ const withBlockTree = ( reducer ) => ( state = {}, action ) => {
 						return result;
 					}, {} ),
 				},
-				updatedBlockUids
+				updatedBlockUids,
+				false
 			);
 		}
 	}