Add insert and remove recursive methods on EntityWorldMut and EntityCommands (#17463)

# Objective

While being able to quickly add / remove components down a tree is
broadly useful (material changing!), it's particularly necessary when
combined with the newly added #13120.

## Solution

Write four methods: covering both adding and removal on both
`EntityWorldMut` and `EntityCommands`.

These methods are generic over the `RelationshipTarget`, thanks to the
freshly merged relations 🎉

## Testing

I've added a simple unit test for these methods.

---------

Co-authored-by: Zachary Harrold <zac@harrold.com.au>
This commit is contained in:
Alice Cecile 2025-01-20 21:57:57 -05:00 committed by GitHub
parent 42b928b90e
commit 85eceb022d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -5,6 +5,7 @@ use crate::{
system::{Commands, EntityCommands},
world::{EntityWorldMut, World},
};
use alloc::vec::Vec;
use core::marker::PhantomData;
impl<'w> EntityWorldMut<'w> {
@ -45,6 +46,55 @@ impl<'w> EntityWorldMut<'w> {
}
self
}
/// Inserts a component or bundle of components into the entity and all related entities,
/// traversing the relationship tracked in `S` in a breadth-first manner.
///
/// # Warning
///
/// This method should only be called on relationships that form a tree-like structure.
/// Any cycles will cause this method to loop infinitely.
// We could keep track of a list of visited entities and track cycles,
// but this is not a very well-defined operation (or hard to write) for arbitrary relationships.
pub fn insert_recursive<S: RelationshipTarget>(
&mut self,
bundle: impl Bundle + Clone,
) -> &mut Self {
self.insert(bundle.clone());
if let Some(relationship_target) = self.get::<S>() {
let related_vec: Vec<Entity> = relationship_target.iter().collect();
for related in related_vec {
self.world_scope(|world| {
world
.entity_mut(related)
.insert_recursive::<S>(bundle.clone());
});
}
}
self
}
/// Removes a component or bundle of components of type `B` from the entity and all related entities,
/// traversing the relationship tracked in `S` in a breadth-first manner.
///
/// # Warning
///
/// This method should only be called on relationships that form a tree-like structure.
/// Any cycles will cause this method to loop infinitely.
pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
self.remove::<B>();
if let Some(relationship_target) = self.get::<S>() {
let related_vec: Vec<Entity> = relationship_target.iter().collect();
for related in related_vec {
self.world_scope(|world| {
world.entity_mut(related).remove_recursive::<S, B>();
});
}
}
self
}
}
impl<'a> EntityCommands<'a> {
@ -79,6 +129,39 @@ impl<'a> EntityCommands<'a> {
});
self
}
/// Inserts a component or bundle of components into the entity and all related entities,
/// traversing the relationship tracked in `S` in a breadth-first manner.
///
/// # Warning
///
/// This method should only be called on relationships that form a tree-like structure.
/// Any cycles will cause this method to loop infinitely.
pub fn insert_recursive<S: RelationshipTarget>(
&mut self,
bundle: impl Bundle + Clone,
) -> &mut Self {
let id = self.id();
self.commands.queue(move |world: &mut World| {
world.entity_mut(id).insert_recursive::<S>(bundle);
});
self
}
/// Removes a component or bundle of components of type `B` from the entity and all related entities,
/// traversing the relationship tracked in `S` in a breadth-first manner.
///
/// # Warning
///
/// This method should only be called on relationships that form a tree-like structure.
/// Any cycles will cause this method to loop infinitely.
pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
let id = self.id();
self.commands.queue(move |world: &mut World| {
world.entity_mut(id).remove_recursive::<S, B>();
});
self
}
}
/// Directly spawns related "source" entities with the given [`Relationship`], targeting
@ -162,3 +245,52 @@ impl<'w, R: Relationship> RelatedSpawnerCommands<'w, R> {
&mut self.commands
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate as bevy_ecs;
use crate::prelude::{ChildOf, Children, Component};
#[derive(Component, Clone, Copy)]
struct TestComponent;
#[test]
fn insert_and_remove_recursive() {
let mut world = World::new();
let a = world.spawn_empty().id();
let b = world.spawn(ChildOf(a)).id();
let c = world.spawn(ChildOf(a)).id();
let d = world.spawn(ChildOf(b)).id();
world
.entity_mut(a)
.insert_recursive::<Children>(TestComponent);
for entity in [a, b, c, d] {
assert!(world.entity(entity).contains::<TestComponent>());
}
world
.entity_mut(b)
.remove_recursive::<Children, TestComponent>();
// Parent
assert!(world.entity(a).contains::<TestComponent>());
// Target
assert!(!world.entity(b).contains::<TestComponent>());
// Sibling
assert!(world.entity(c).contains::<TestComponent>());
// Child
assert!(!world.entity(d).contains::<TestComponent>());
world
.entity_mut(a)
.remove_recursive::<Children, TestComponent>();
for entity in [a, b, c, d] {
assert!(!world.entity(entity).contains::<TestComponent>());
}
}
}