Skip to content

Commit 85eceb0

Browse files
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 <[email protected]>
1 parent 42b928b commit 85eceb0

File tree

1 file changed

+132
-0
lines changed

1 file changed

+132
-0
lines changed

crates/bevy_ecs/src/relationship/related_methods.rs

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use crate::{
55
system::{Commands, EntityCommands},
66
world::{EntityWorldMut, World},
77
};
8+
use alloc::vec::Vec;
89
use core::marker::PhantomData;
910

1011
impl<'w> EntityWorldMut<'w> {
@@ -45,6 +46,55 @@ impl<'w> EntityWorldMut<'w> {
4546
}
4647
self
4748
}
49+
50+
/// Inserts a component or bundle of components into the entity and all related entities,
51+
/// traversing the relationship tracked in `S` in a breadth-first manner.
52+
///
53+
/// # Warning
54+
///
55+
/// This method should only be called on relationships that form a tree-like structure.
56+
/// Any cycles will cause this method to loop infinitely.
57+
// We could keep track of a list of visited entities and track cycles,
58+
// but this is not a very well-defined operation (or hard to write) for arbitrary relationships.
59+
pub fn insert_recursive<S: RelationshipTarget>(
60+
&mut self,
61+
bundle: impl Bundle + Clone,
62+
) -> &mut Self {
63+
self.insert(bundle.clone());
64+
if let Some(relationship_target) = self.get::<S>() {
65+
let related_vec: Vec<Entity> = relationship_target.iter().collect();
66+
for related in related_vec {
67+
self.world_scope(|world| {
68+
world
69+
.entity_mut(related)
70+
.insert_recursive::<S>(bundle.clone());
71+
});
72+
}
73+
}
74+
75+
self
76+
}
77+
78+
/// Removes a component or bundle of components of type `B` from the entity and all related entities,
79+
/// traversing the relationship tracked in `S` in a breadth-first manner.
80+
///
81+
/// # Warning
82+
///
83+
/// This method should only be called on relationships that form a tree-like structure.
84+
/// Any cycles will cause this method to loop infinitely.
85+
pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
86+
self.remove::<B>();
87+
if let Some(relationship_target) = self.get::<S>() {
88+
let related_vec: Vec<Entity> = relationship_target.iter().collect();
89+
for related in related_vec {
90+
self.world_scope(|world| {
91+
world.entity_mut(related).remove_recursive::<S, B>();
92+
});
93+
}
94+
}
95+
96+
self
97+
}
4898
}
4999

50100
impl<'a> EntityCommands<'a> {
@@ -79,6 +129,39 @@ impl<'a> EntityCommands<'a> {
79129
});
80130
self
81131
}
132+
133+
/// Inserts a component or bundle of components into the entity and all related entities,
134+
/// traversing the relationship tracked in `S` in a breadth-first manner.
135+
///
136+
/// # Warning
137+
///
138+
/// This method should only be called on relationships that form a tree-like structure.
139+
/// Any cycles will cause this method to loop infinitely.
140+
pub fn insert_recursive<S: RelationshipTarget>(
141+
&mut self,
142+
bundle: impl Bundle + Clone,
143+
) -> &mut Self {
144+
let id = self.id();
145+
self.commands.queue(move |world: &mut World| {
146+
world.entity_mut(id).insert_recursive::<S>(bundle);
147+
});
148+
self
149+
}
150+
151+
/// Removes a component or bundle of components of type `B` from the entity and all related entities,
152+
/// traversing the relationship tracked in `S` in a breadth-first manner.
153+
///
154+
/// # Warning
155+
///
156+
/// This method should only be called on relationships that form a tree-like structure.
157+
/// Any cycles will cause this method to loop infinitely.
158+
pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
159+
let id = self.id();
160+
self.commands.queue(move |world: &mut World| {
161+
world.entity_mut(id).remove_recursive::<S, B>();
162+
});
163+
self
164+
}
82165
}
83166

84167
/// Directly spawns related "source" entities with the given [`Relationship`], targeting
@@ -162,3 +245,52 @@ impl<'w, R: Relationship> RelatedSpawnerCommands<'w, R> {
162245
&mut self.commands
163246
}
164247
}
248+
249+
#[cfg(test)]
250+
mod tests {
251+
use super::*;
252+
use crate as bevy_ecs;
253+
use crate::prelude::{ChildOf, Children, Component};
254+
255+
#[derive(Component, Clone, Copy)]
256+
struct TestComponent;
257+
258+
#[test]
259+
fn insert_and_remove_recursive() {
260+
let mut world = World::new();
261+
262+
let a = world.spawn_empty().id();
263+
let b = world.spawn(ChildOf(a)).id();
264+
let c = world.spawn(ChildOf(a)).id();
265+
let d = world.spawn(ChildOf(b)).id();
266+
267+
world
268+
.entity_mut(a)
269+
.insert_recursive::<Children>(TestComponent);
270+
271+
for entity in [a, b, c, d] {
272+
assert!(world.entity(entity).contains::<TestComponent>());
273+
}
274+
275+
world
276+
.entity_mut(b)
277+
.remove_recursive::<Children, TestComponent>();
278+
279+
// Parent
280+
assert!(world.entity(a).contains::<TestComponent>());
281+
// Target
282+
assert!(!world.entity(b).contains::<TestComponent>());
283+
// Sibling
284+
assert!(world.entity(c).contains::<TestComponent>());
285+
// Child
286+
assert!(!world.entity(d).contains::<TestComponent>());
287+
288+
world
289+
.entity_mut(a)
290+
.remove_recursive::<Children, TestComponent>();
291+
292+
for entity in [a, b, c, d] {
293+
assert!(!world.entity(entity).contains::<TestComponent>());
294+
}
295+
}
296+
}

0 commit comments

Comments
 (0)