Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Taffy viewport node leaks #17596

Merged
merged 6 commits into from
Feb 2, 2025
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Changed the UiSurface::entity_to_taffy to map to LayoutNodes.
`LayoutNode`s also have a `viewport_id: Option<taffy::NodeId`  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.
  • Loading branch information
ickshonpe committed Jan 29, 2025
commit 1f17ff918ac6d560212d77501cb8bf325e55c0a8
2 changes: 1 addition & 1 deletion crates/bevy_ui/src/layout/debug.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
29 changes: 16 additions & 13 deletions crates/bevy_ui/src/layout/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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(|_| {
Expand All @@ -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()
);

Expand All @@ -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
Expand All @@ -737,20 +737,23 @@ 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()
);

// the remaining children should still have nodes in the layout tree
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
Expand All @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -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;
Expand Down Expand Up @@ -1244,6 +1247,6 @@ 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());
}
}
114 changes: 73 additions & 41 deletions crates/bevy_ui/src/layout/ui_surface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,25 @@ 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>,
Expand Down Expand Up @@ -77,70 +93,75 @@ 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,
)
} else {
taffy.new_leaf(convert::from_node(node, layout_context, false))
};
entry.insert(taffy_node_id.unwrap());
entry.insert(taffy_node.unwrap().into());
}
}
}

/// 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`].
pub fn update_children(&mut self, entity: Entity, children: impl Iterator<Item = Entity>) {
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).unwrap();
}
}
}

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();
}
}

Expand All @@ -167,25 +188,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);
Expand Down Expand Up @@ -265,11 +287,15 @@ impl UiSurface {
}
}

/// 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();
println!("removed {:?}", node);
self.taffy.remove(node.id).unwrap();
if let Some(viewport_node) = node.viewport_id {
self.taffy.remove(viewport_node).unwrap();
}
}
}
}
Expand All @@ -293,10 +319,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))
}
Expand Down Expand Up @@ -353,7 +379,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]
Expand Down Expand Up @@ -455,7 +481,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!(
Expand Down Expand Up @@ -596,7 +625,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(
Expand All @@ -620,7 +649,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());
Expand All @@ -646,14 +675,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"
);
Expand All @@ -680,13 +712,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"
);
Expand All @@ -712,13 +744,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"
);
Expand Down
Loading