forestia/bevy
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix Taffy viewport node leaks (#17596)

Browse Source 
# 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::<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
        );
    }
```

Fixes #17594

## Solution

Change the `UiSurface::entity_to_taffy` to map to `LayoutNode`s. A
`LayoutNode` has a `viewport_id: Option<taffy::NodeId>` 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`
This commit is contained in:
ickshonpe 2025-02-02 15:03:10 +00:00 committed by GitHub
parent afef7d5797
commit 89a1c49377
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 156 additions and 56 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
crates/bevy_ui/src/layout/debug.rs
View File

@ -12,7 +12,7 @@ pub fn print_ui_layout_tree(ui_surface: &UiSurface) {
let taffy_to_entity: HashMap<NodeId, Entity> = 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();

93
crates/bevy_ui/src/layout/mod.rs
View File

@ -671,7 +671,7 @@ mod tests {
let ui_surface = world.resource::<UiSurface>();
// `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::<UiSurface>();
// 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::<UiSurface>();
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::<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
);
}
}

117
crates/bevy_ui/src/layout/ui_surface.rs
View File

@ -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<taffy::NodeId>,
// The id of the node in the taffy tree
pub(super) id: taffy::NodeId,
}
impl From<taffy::NodeId> 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<taffy::NodeId>,
pub(super) entity_to_taffy: EntityHashMap<LayoutNode>,
pub(super) camera_entity_to_taffy: EntityHashMap<EntityHashMap<taffy::NodeId>>,
pub(super) camera_roots: EntityHashMap<Vec<RootNodePair>>,
pub(super) taffy: TaffyTree<NodeMeasure>,
@ -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<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;
}
}
}
}
/// 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<Item = Entity>) {
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"
);