From a70e09015696e630ca66c6f341ed212d19ea0bef Mon Sep 17 00:00:00 2001
From: Gabriel Bourgeois <gabriel.bourgeoisv4si@gmail.com>
Date: Mon, 5 Sep 2022 21:50:31 +0000
Subject: [PATCH] Clean up taffy nodes when UI node entities are removed
 (#5886)

# Objective

Clean up taffy nodes when the associated UI node gets removed. The current UI code will keep the taffy nodes around forever.

## Solution

Use `RemovedComponents<Node>` to iterate over nodes that are no longer valid UI nodes or that have been despawned, and remove them from taffy and the internal hash map.

## Implementation Notes

Do note that using `despawn()` instead of `despawn_recursive()` on a UI node that has children will result in a [warnings spam](https://github.com/bevyengine/bevy/blob/main/crates/bevy_ui/src/flex/mod.rs#L120) since the children will not be part of a proper UI hierarchy anymore.

---

## Changelog

- Fixed memory leak when nodes are removed in bevy_ui
---
 crates/bevy_ecs/src/system/system_param.rs |  9 +++++++++
 crates/bevy_ui/src/flex/mod.rs             | 15 +++++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs
index 8762bd64c4255..fa6ad3f26fb7b 100644
--- a/crates/bevy_ecs/src/system/system_param.rs
+++ b/crates/bevy_ecs/src/system/system_param.rs
@@ -780,6 +780,15 @@ impl<'a, T: Component> RemovedComponents<'a, T> {
     }
 }
 
+impl<'a, T: Component> IntoIterator for &'a RemovedComponents<'a, T> {
+    type Item = Entity;
+    type IntoIter = std::iter::Cloned<std::slice::Iter<'a, Entity>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
 // SAFETY: Only reads World components
 unsafe impl<T: Component> ReadOnlySystemParamFetch for RemovedComponentsState<T> {}
 
diff --git a/crates/bevy_ui/src/flex/mod.rs b/crates/bevy_ui/src/flex/mod.rs
index 9e98e1f01d7f5..633d45aff5919 100644
--- a/crates/bevy_ui/src/flex/mod.rs
+++ b/crates/bevy_ui/src/flex/mod.rs
@@ -5,7 +5,7 @@ use bevy_ecs::{
     entity::Entity,
     event::EventReader,
     query::{Changed, With, Without, WorldQuery},
-    system::{Query, Res, ResMut, Resource},
+    system::{Query, RemovedComponents, Res, ResMut, Resource},
 };
 use bevy_hierarchy::{Children, Parent};
 use bevy_log::warn;
@@ -172,6 +172,15 @@ without UI components as a child of an entity with UI components, results may be
         }
     }
 
+    /// Removes each entity from the internal map and then removes their associated node from taffy
+    pub fn remove_entities(&mut self, entities: impl IntoIterator<Item = Entity>) {
+        for entity in entities {
+            if let Some(node) = self.entity_to_taffy.remove(&entity) {
+                self.taffy.remove(node);
+            }
+        }
+    }
+
     pub fn get_layout(&self, entity: Entity) -> Result<&taffy::layout::Layout, FlexError> {
         if let Some(taffy_node) = self.entity_to_taffy.get(&entity) {
             self.taffy
@@ -208,6 +217,7 @@ pub fn flex_node_system(
     >,
     children_query: Query<(Entity, &Children), (With<Node>, Changed<Children>)>,
     mut node_transform_query: Query<(Entity, &mut Node, &mut Transform, Option<&Parent>)>,
+    removed_nodes: RemovedComponents<Node>,
 ) {
     // update window root nodes
     for window in windows.iter() {
@@ -244,7 +254,8 @@ pub fn flex_node_system(
         flex_surface.upsert_leaf(entity, style, *calculated_size, scale_factor);
     }
 
-    // TODO: handle removed nodes
+    // clean up removed nodes
+    flex_surface.remove_entities(&removed_nodes);
 
     // update window children (for now assuming all Nodes live in the primary window)
     if let Some(primary_window) = windows.get_primary() {