From 536133f25af3db8eda4aed22f11278054a4ab9ff Mon Sep 17 00:00:00 2001
From: ickshonpe <david.curthoys@googlemail.com>
Date: Thu, 30 Jan 2025 12:58:06 +0000
Subject: [PATCH] Clear `viewport_id` from `entity_to_taffy` in
 `remove_camera_entities`. Add tests to check for leaked vewport nodes.

---
 crates/bevy_ui/src/layout/mod.rs        | 64 +++++++++++++++++++++++++
 crates/bevy_ui/src/layout/ui_surface.rs | 11 ++---
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_ui/src/layout/mod.rs b/crates/bevy_ui/src/layout/mod.rs
index 5c2af9e1a5216..eab2c92a5c6d5 100644
--- a/crates/bevy_ui/src/layout/mod.rs
+++ b/crates/bevy_ui/src/layout/mod.rs
@@ -1249,4 +1249,68 @@ mod tests {
         let taffy_node = ui_surface.entity_to_taffy.get(&root_node_entity).unwrap();
         assert!(ui_surface.taffy.layout(taffy_node.id).is_ok());
     }
+
+    #[test]
+    fn no_viewport_node_leak_on_root_despawned() {
+        let (mut world, mut ui_schedule) = setup_ui_test_world();
+
+        let ui_root_entity = world.spawn(Node::default()).id();
+
+        // The UI schedule synchronises Bevy UI's internal `TaffyTree` with the
+        // main world's tree of `Node` entities.
+        ui_schedule.run(&mut world);
+
+        // Two taffy nodes are added to the internal `TaffyTree` for each root UI entity.
+        // An implicit taffy node representing the viewport and a taffy node corresponding to the
+        // root UI entity which is parented to the viewport taffy node.
+        assert_eq!(
+            world.resource_mut::<UiSurface>().taffy.total_node_count(),
+            2
+        );
+
+        world.despawn(ui_root_entity);
+
+        // The UI schedule removes both the taffy node corresponding to `ui_root_entity` and its
+        // parent viewport node.
+        ui_schedule.run(&mut world);
+
+        // Both taffy nodes should now be removed from the internal `TaffyTree`
+        assert_eq!(
+            world.resource_mut::<UiSurface>().taffy.total_node_count(),
+            0
+        );
+    }
+
+    #[test]
+    fn no_viewport_node_leak_on_parented_root() {
+        let (mut world, mut ui_schedule) = setup_ui_test_world();
+
+        let ui_root_entity_1 = world.spawn(Node::default()).id();
+        let ui_root_entity_2 = world.spawn(Node::default()).id();
+
+        ui_schedule.run(&mut world);
+
+        // There are two UI root entities. Each root taffy node is given it's own viewport node parent,
+        // so a total of four taffy nodes are added to the `TaffyTree` by the UI schedule.
+        assert_eq!(
+            world.resource_mut::<UiSurface>().taffy.total_node_count(),
+            4
+        );
+
+        // Parent `ui_root_entity_2` onto `ui_root_entity_1` so now only `ui_root_entity_1` is a
+        // UI root entity.
+        world
+            .entity_mut(ui_root_entity_1)
+            .add_child(ui_root_entity_2);
+
+        // Now there is only one root node so the second viewport node is removed by
+        // the UI schedule.
+        ui_schedule.run(&mut world);
+
+        // There is only one viewport node now, so the `TaffyTree` contains 3 nodes in total.
+        assert_eq!(
+            world.resource_mut::<UiSurface>().taffy.total_node_count(),
+            3
+        );
+    }
 }
diff --git a/crates/bevy_ui/src/layout/ui_surface.rs b/crates/bevy_ui/src/layout/ui_surface.rs
index f9fd3e3ebbf72..661924229efd8 100644
--- a/crates/bevy_ui/src/layout/ui_surface.rs
+++ b/crates/bevy_ui/src/layout/ui_surface.rs
@@ -141,9 +141,7 @@ impl UiSurface {
             if let Some(taffy_node) = self.entity_to_taffy.get_mut(&child) {
                 self.taffy_children_scratch.push(taffy_node.id);
                 if let Some(viewport_id) = taffy_node.viewport_id.take() {
-                    if self.taffy.get_node_context(viewport_id).is_some() {
-                        self.taffy.remove(viewport_id).ok();
-                    }
+                    self.taffy.remove(viewport_id).ok();
                 }
             }
         }
@@ -283,8 +281,9 @@ impl UiSurface {
     pub fn remove_camera_entities(&mut self, entities: impl IntoIterator<Item = Entity>) {
         for entity in entities {
             if let Some(camera_root_node_map) = self.camera_entity_to_taffy.remove(&entity) {
-                for (_, node) in camera_root_node_map.iter() {
+                for (entity, node) in camera_root_node_map.iter() {
                     self.taffy.remove(*node).unwrap();
+                    self.entity_to_taffy.get_mut(entity).unwrap().viewport_id = None;
                 }
             }
         }
@@ -296,9 +295,7 @@ impl UiSurface {
             if let Some(node) = self.entity_to_taffy.remove(&entity) {
                 self.taffy.remove(node.id).unwrap();
                 if let Some(viewport_node) = node.viewport_id {
-                    if self.taffy.get_node_context(viewport_node).is_some() {
-                        self.taffy.remove(viewport_node).ok();
-                    }
+                    self.taffy.remove(viewport_node).ok();
                 }
             }
         }