From 89a1c493771dd9bbf3ad23b3897b89e58ff45934 Mon Sep 17 00:00:00 2001 From: ickshonpe Date: Sun, 2 Feb 2025 15:03:10 +0000 Subject: [PATCH] Fix Taffy viewport node leaks (#17596) # Objective For most UI node entities there's a 1-to-1 mapping from the entity to its associated Taffy node. Root UI nodes are an exception though, their corresponding Taffy node in the Taffy tree is also given a parent that represents the viewport. These viewport Taffy nodes are not removed when a root UI node is despawned. Parenting of an existing root UI node with an associated viewport Taffy node also results in the leak of the viewport node. These tests fail if added to the `layout` module's tests on the main branch: ```rust #[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 synchronizes 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::().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::().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::().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::().taffy.total_node_count(), 3 ); } ``` Fixes #17594 ## Solution Change the `UiSurface::entity_to_taffy` to map to `LayoutNode`s. A `LayoutNode` has a `viewport_id: Option` field which is the id of the corresponding implicit "viewport" node if the node is a root UI node, otherwise it is `None`. When removing or parenting nodes this field is checked and the implicit viewport node is removed if present. ## Testing There are two new tests in `bevy_ui::layout::tests` included with this PR: * `no_viewport_node_leak_on_root_despawned` * `no_viewport_node_leak_on_parented_root` --- crates/bevy_ui/src/layout/debug.rs | 2 +- crates/bevy_ui/src/layout/mod.rs | 93 ++++++++++++++++--- crates/bevy_ui/src/layout/ui_surface.rs | 117 +++++++++++++++--------- 3 files changed, 156 insertions(+), 56 deletions(-) diff --git a/crates/bevy_ui/src/layout/debug.rs b/crates/bevy_ui/src/layout/debug.rs index 5af75d3ffdc32..e83f2a3b23e99 100644 --- a/crates/bevy_ui/src/layout/debug.rs +++ b/crates/bevy_ui/src/layout/debug.rs @@ -12,7 +12,7 @@ pub fn print_ui_layout_tree(ui_surface: &UiSurface) { let taffy_to_entity: HashMap = ui_surface .entity_to_taffy .iter() - .map(|(entity, node)| (*node, *entity)) + .map(|(entity, node)| (node.id, *entity)) .collect(); for (&entity, roots) in &ui_surface.camera_roots { let mut out = String::new(); diff --git a/crates/bevy_ui/src/layout/mod.rs b/crates/bevy_ui/src/layout/mod.rs index 7288fbec1bdef..6f7875e62a7e0 100644 --- a/crates/bevy_ui/src/layout/mod.rs +++ b/crates/bevy_ui/src/layout/mod.rs @@ -671,7 +671,7 @@ mod tests { let ui_surface = world.resource::(); // `ui_node` is removed, attempting to retrieve a style for `ui_node` panics - let _ = ui_surface.taffy.style(ui_node); + let _ = ui_surface.taffy.style(ui_node.id); } #[test] @@ -687,7 +687,7 @@ mod tests { let ui_parent_node = ui_surface.entity_to_taffy[&ui_parent_entity]; // `ui_parent_node` shouldn't have any children yet - assert_eq!(ui_surface.taffy.child_count(ui_parent_node), 0); + assert_eq!(ui_surface.taffy.child_count(ui_parent_node.id), 0); let mut ui_child_entities = (0..10) .map(|_| { @@ -706,7 +706,7 @@ mod tests { 1 + ui_child_entities.len() ); assert_eq!( - ui_surface.taffy.child_count(ui_parent_node), + ui_surface.taffy.child_count(ui_parent_node.id), ui_child_entities.len() ); @@ -718,7 +718,7 @@ mod tests { // the children should have a corresponding ui node and that ui node's parent should be `ui_parent_node` for node in child_node_map.values() { - assert_eq!(ui_surface.taffy.parent(*node), Some(ui_parent_node)); + assert_eq!(ui_surface.taffy.parent(node.id), Some(ui_parent_node.id)); } // delete every second child @@ -737,7 +737,7 @@ mod tests { 1 + ui_child_entities.len() ); assert_eq!( - ui_surface.taffy.child_count(ui_parent_node), + ui_surface.taffy.child_count(ui_parent_node.id), ui_child_entities.len() ); @@ -745,12 +745,15 @@ mod tests { for child_entity in &ui_child_entities { let child_node = child_node_map[child_entity]; assert_eq!(ui_surface.entity_to_taffy[child_entity], child_node); - assert_eq!(ui_surface.taffy.parent(child_node), Some(ui_parent_node)); + assert_eq!( + ui_surface.taffy.parent(child_node.id), + Some(ui_parent_node.id) + ); assert!(ui_surface .taffy - .children(ui_parent_node) + .children(ui_parent_node.id) .unwrap() - .contains(&child_node)); + .contains(&child_node.id)); } // the nodes of the deleted children should have been removed from the layout tree @@ -761,9 +764,9 @@ mod tests { let deleted_child_node = child_node_map[deleted_child_entity]; assert!(!ui_surface .taffy - .children(ui_parent_node) + .children(ui_parent_node.id) .unwrap() - .contains(&deleted_child_node)); + .contains(&deleted_child_node.id)); } // despawn the parent entity and its descendants @@ -1069,7 +1072,7 @@ mod tests { let ui_node = ui_surface.entity_to_taffy[&ui_entity]; // a node with a content size should have taffy context - assert!(ui_surface.taffy.get_node_context(ui_node).is_some()); + assert!(ui_surface.taffy.get_node_context(ui_node.id).is_some()); let layout = ui_surface.get_layout(ui_entity, true).unwrap().0; assert_eq!(layout.size.width, content_size.x); assert_eq!(layout.size.height, content_size.y); @@ -1080,7 +1083,7 @@ mod tests { let mut ui_surface = world.resource_mut::(); // a node without a content size should not have taffy context - assert!(ui_surface.taffy.get_node_context(ui_node).is_none()); + assert!(ui_surface.taffy.get_node_context(ui_node.id).is_none()); // Without a content size, the node has no width or height constraints so the length of both dimensions is 0. let layout = ui_surface.get_layout(ui_entity, true).unwrap().0; @@ -1244,6 +1247,70 @@ mod tests { let ui_surface = world.resource::(); let taffy_node = ui_surface.entity_to_taffy.get(&root_node_entity).unwrap(); - assert!(ui_surface.taffy.layout(*taffy_node).is_ok()); + 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 synchronizes 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::().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::().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::().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::().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 92bbed4a7177d..661924229efd8 100644 --- a/crates/bevy_ui/src/layout/ui_surface.rs +++ b/crates/bevy_ui/src/layout/ui_surface.rs @@ -21,9 +21,26 @@ pub struct RootNodePair { pub(super) user_root_node: taffy::NodeId, } +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct LayoutNode { + // Implicit "viewport" node if this `LayoutNode` corresponds to a root UI node entity + pub(super) viewport_id: Option, + // The id of the node in the taffy tree + pub(super) id: taffy::NodeId, +} + +impl From for LayoutNode { + fn from(value: taffy::NodeId) -> Self { + LayoutNode { + viewport_id: None, + id: value, + } + } +} + #[derive(Resource)] pub struct UiSurface { - pub(super) entity_to_taffy: EntityHashMap, + pub(super) entity_to_taffy: EntityHashMap, pub(super) camera_entity_to_taffy: EntityHashMap>, pub(super) camera_roots: EntityHashMap>, pub(super) taffy: TaffyTree, @@ -77,25 +94,25 @@ impl UiSurface { match self.entity_to_taffy.entry(entity) { Entry::Occupied(entry) => { - let taffy_node_id = *entry.get(); + let taffy_node = *entry.get(); let has_measure = if new_node_context.is_some() { taffy - .set_node_context(taffy_node_id, new_node_context) + .set_node_context(taffy_node.id, new_node_context) .unwrap(); true } else { - taffy.get_node_context(taffy_node_id).is_some() + taffy.get_node_context(taffy_node.id).is_some() }; taffy .set_style( - taffy_node_id, + taffy_node.id, convert::from_node(node, layout_context, has_measure), ) .unwrap(); } Entry::Vacant(entry) => { - let taffy_node_id = if let Some(measure) = new_node_context.take() { + let taffy_node = if let Some(measure) = new_node_context.take() { taffy.new_leaf_with_context( convert::from_node(node, layout_context, true), measure, @@ -103,7 +120,7 @@ impl UiSurface { } else { taffy.new_leaf(convert::from_node(node, layout_context, false)) }; - entry.insert(taffy_node_id.unwrap()); + entry.insert(taffy_node.unwrap().into()); } } } @@ -111,7 +128,9 @@ impl UiSurface { /// Update the `MeasureFunc` of the taffy node corresponding to the given [`Entity`] if the node exists. pub fn update_node_context(&mut self, entity: Entity, context: NodeMeasure) -> Option<()> { let taffy_node = self.entity_to_taffy.get(&entity)?; - self.taffy.set_node_context(*taffy_node, Some(context)).ok() + self.taffy + .set_node_context(taffy_node.id, Some(context)) + .ok() } /// Update the children of the taffy node corresponding to the given [`Entity`]. @@ -119,28 +138,31 @@ impl UiSurface { self.taffy_children_scratch.clear(); for child in children { - if let Some(taffy_node) = self.entity_to_taffy.get(&child) { - self.taffy_children_scratch.push(*taffy_node); + 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() { + self.taffy.remove(viewport_id).ok(); + } } } let taffy_node = self.entity_to_taffy.get(&entity).unwrap(); self.taffy - .set_children(*taffy_node, &self.taffy_children_scratch) + .set_children(taffy_node.id, &self.taffy_children_scratch) .unwrap(); } /// Removes children from the entity's taffy node if it exists. Does nothing otherwise. pub fn try_remove_children(&mut self, entity: Entity) { if let Some(taffy_node) = self.entity_to_taffy.get(&entity) { - self.taffy.set_children(*taffy_node, &[]).unwrap(); + self.taffy.set_children(taffy_node.id, &[]).unwrap(); } } /// Removes the measure from the entity's taffy node if it exists. Does nothing otherwise. pub fn try_remove_node_context(&mut self, entity: Entity) { if let Some(taffy_node) = self.entity_to_taffy.get(&entity) { - self.taffy.set_node_context(*taffy_node, None).unwrap(); + self.taffy.set_node_context(taffy_node.id, None).unwrap(); } } @@ -167,25 +189,26 @@ impl UiSurface { let existing_roots = self.camera_roots.entry(camera_id).or_default(); let mut new_roots = Vec::new(); for entity in children { - let node = *self.entity_to_taffy.get(&entity).unwrap(); + let node = self.entity_to_taffy.get_mut(&entity).unwrap(); let root_node = existing_roots .iter() - .find(|n| n.user_root_node == node) + .find(|n| n.user_root_node == node.id) .cloned() .unwrap_or_else(|| { - if let Some(previous_parent) = self.taffy.parent(node) { + if let Some(previous_parent) = self.taffy.parent(node.id) { // remove the root node from the previous implicit node's children - self.taffy.remove_child(previous_parent, node).unwrap(); + self.taffy.remove_child(previous_parent, node.id).unwrap(); } - let viewport_node = *camera_root_node_map - .entry(entity) - .or_insert_with(|| self.taffy.new_leaf(viewport_style.clone()).unwrap()); - self.taffy.add_child(viewport_node, node).unwrap(); - + let viewport_node = *camera_root_node_map.entry(entity).or_insert_with(|| { + node.viewport_id + .unwrap_or_else(|| self.taffy.new_leaf(viewport_style.clone()).unwrap()) + }); + node.viewport_id = Some(viewport_node); + self.taffy.add_child(viewport_node, node.id).unwrap(); RootNodePair { implicit_viewport_node: viewport_node, - user_root_node: node, + user_root_node: node.id, } }); new_roots.push(root_node); @@ -258,18 +281,22 @@ impl UiSurface { pub fn remove_camera_entities(&mut self, entities: impl IntoIterator) { 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; } } } } - /// Removes each entity from the internal map and then removes their associated node from taffy + /// Removes each entity from the internal map and then removes their associated nodes from taffy pub fn remove_entities(&mut self, entities: impl IntoIterator) { for entity in entities { if let Some(node) = self.entity_to_taffy.remove(&entity) { - self.taffy.remove(node).unwrap(); + self.taffy.remove(node.id).unwrap(); + if let Some(viewport_node) = node.viewport_id { + self.taffy.remove(viewport_node).ok(); + } } } } @@ -293,10 +320,10 @@ impl UiSurface { self.taffy.disable_rounding(); } - let out = match self.taffy.layout(*taffy_node).cloned() { + let out = match self.taffy.layout(taffy_node.id).cloned() { Ok(layout) => { self.taffy.disable_rounding(); - let taffy_size = self.taffy.layout(*taffy_node).unwrap().size; + let taffy_size = self.taffy.layout(taffy_node.id).unwrap().size; let unrounded_size = Vec2::new(taffy_size.width, taffy_size.height); Ok((layout, unrounded_size)) } @@ -353,7 +380,7 @@ mod tests { let root_node_taffy = ui_surface.entity_to_taffy.get(&root_node_entity)?; root_node_pairs .iter() - .find(|&root_node_pair| root_node_pair.user_root_node == *root_node_taffy) + .find(|&root_node_pair| root_node_pair.user_root_node == root_node_taffy.id) } #[test] @@ -455,7 +482,10 @@ mod tests { assert!(root_node_pair.is_some()); assert_eq!( Some(root_node_pair.unwrap().user_root_node).as_ref(), - ui_surface.entity_to_taffy.get(&root_node_entity) + ui_surface + .entity_to_taffy + .get(&root_node_entity) + .map(|taffy_node| &taffy_node.id) ); assert_eq!( @@ -596,7 +626,7 @@ mod tests { let parent_node = *ui_surface.entity_to_taffy.get(&root_node_entity).unwrap(); let child_node = *ui_surface.entity_to_taffy.get(&child_entity).unwrap(); - assert_eq!(ui_surface.taffy.parent(child_node), Some(parent_node)); + assert_eq!(ui_surface.taffy.parent(child_node.id), Some(parent_node.id)); } #[expect( @@ -620,7 +650,7 @@ mod tests { // set up the relationship manually ui_surface .taffy - .add_child(root_taffy_node, child_taffy) + .add_child(root_taffy_node.id, child_taffy.id) .unwrap(); ui_surface.set_camera_children(camera_entity, [root_node_entity].into_iter()); @@ -646,14 +676,17 @@ mod tests { get_root_node_pair_exact(&ui_surface, root_node_entity, camera_entity) .expect("expected root node pair"); - assert_eq!(ui_surface.taffy.parent(child_taffy), Some(root_taffy_node)); - let root_taffy_children = ui_surface.taffy.children(root_taffy_node).unwrap(); + assert_eq!( + ui_surface.taffy.parent(child_taffy.id), + Some(root_taffy_node.id) + ); + let root_taffy_children = ui_surface.taffy.children(root_taffy_node.id).unwrap(); assert!( - root_taffy_children.contains(&child_taffy), + root_taffy_children.contains(&child_taffy.id), "root node is not a parent of child node" ); assert_eq!( - ui_surface.taffy.child_count(root_taffy_node), + ui_surface.taffy.child_count(root_taffy_node.id), 1, "expected root node child count to be 1" ); @@ -680,13 +713,13 @@ mod tests { "child of root node should not be associated with camera" ); - let root_taffy_children = ui_surface.taffy.children(root_taffy_node).unwrap(); + let root_taffy_children = ui_surface.taffy.children(root_taffy_node.id).unwrap(); assert!( - root_taffy_children.contains(&child_taffy), + root_taffy_children.contains(&child_taffy.id), "root node is not a parent of child node" ); assert_eq!( - ui_surface.taffy.child_count(root_taffy_node), + ui_surface.taffy.child_count(root_taffy_node.id), 1, "expected root node child count to be 1" ); @@ -712,13 +745,13 @@ mod tests { ); let child_taffy = ui_surface.entity_to_taffy.get(&child_entity).unwrap(); - let root_taffy_children = ui_surface.taffy.children(root_taffy_node).unwrap(); + let root_taffy_children = ui_surface.taffy.children(root_taffy_node.id).unwrap(); assert!( - root_taffy_children.contains(child_taffy), + root_taffy_children.contains(&child_taffy.id), "root node is not a parent of child node" ); assert_eq!( - ui_surface.taffy.child_count(root_taffy_node), + ui_surface.taffy.child_count(root_taffy_node.id), 1, "expected root node child count to be 1" );