From 6e67d3e534296f9a022b9594f8a3c48b25a50af1 Mon Sep 17 00:00:00 2001
From: ickshonpe <david.curthoys@googlemail.com>
Date: Wed, 5 Apr 2023 22:29:29 +0100
Subject: [PATCH] Remove the corresponding measure from Taffy when a
 `CalculatedSize` component is removed. (#8294)

# Objective

When a `CalculatedSize` component from a UI Node entity is removed, the
corresponding Taffy measure isn't removed which will mess up the layout
in confusing, unpredictable ways.

## Solution

Iterate through all the entities with removed `CalculatedSize`
components and remove the corresponding Taffy measures.
---
 crates/bevy_ui/src/flex/mod.rs | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/crates/bevy_ui/src/flex/mod.rs b/crates/bevy_ui/src/flex/mod.rs
index 62ab50196e..d59d8c5441 100644
--- a/crates/bevy_ui/src/flex/mod.rs
+++ b/crates/bevy_ui/src/flex/mod.rs
@@ -166,6 +166,13 @@ without UI components as a child of an entity with UI components, results may be
         }
     }
 
+    /// Removes the measure from the entity's taffy node if it exists. Does nothing otherwise.
+    pub fn try_remove_measure(&mut self, entity: Entity) {
+        if let Some(taffy_node) = self.entity_to_taffy.get(&entity) {
+            self.taffy.set_measure(*taffy_node, None).unwrap();
+        }
+    }
+
     pub fn update_window(&mut self, window: Entity, window_resolution: &WindowResolution) {
         let taffy = &mut self.taffy;
         let node = self
@@ -258,6 +265,7 @@ pub fn flex_node_system(
     >,
     children_query: Query<(Entity, &Children), (With<Node>, Changed<Children>)>,
     mut removed_children: RemovedComponents<Children>,
+    mut removed_calculated_sizes: RemovedComponents<CalculatedSize>,
     mut node_transform_query: Query<(Entity, &mut Node, &mut Transform, Option<&Parent>)>,
     mut removed_nodes: RemovedComponents<Node>,
 ) {
@@ -320,6 +328,11 @@ pub fn flex_node_system(
     // clean up removed nodes
     flex_surface.remove_entities(removed_nodes.iter());
 
+    // When a `CalculatedSize` component is removed from an entity, we need to remove the measure from the corresponding taffy node.
+    for entity in removed_calculated_sizes.iter() {
+        flex_surface.try_remove_measure(entity);
+    }
+
     // update window children (for now assuming all Nodes live in the primary window)
     flex_surface.set_window_children(primary_window_entity, root_node_query.iter());