bevy_ecs/relationship/
related_methods.rs

1use crate::{
2    bundle::Bundle,
3    entity::{hash_set::EntityHashSet, Entity},
4    prelude::Children,
5    relationship::{
6        Relationship, RelationshipHookMode, RelationshipSourceCollection, RelationshipTarget,
7    },
8    system::{Commands, EntityCommands},
9    world::{DeferredWorld, EntityWorldMut, World},
10};
11use bevy_platform::prelude::{Box, Vec};
12use core::{marker::PhantomData, mem};
13
14use super::OrderedRelationshipSourceCollection;
15
16impl<'w> EntityWorldMut<'w> {
17    /// Spawns a entity related to this entity (with the `R` relationship) by taking a bundle
18    pub fn with_related<R: Relationship>(&mut self, bundle: impl Bundle) -> &mut Self {
19        let parent = self.id();
20        self.world_scope(|world| {
21            world.spawn((bundle, R::from(parent)));
22        });
23        self
24    }
25
26    /// Spawns entities related to this entity (with the `R` relationship) by taking a function that operates on a [`RelatedSpawner`].
27    pub fn with_related_entities<R: Relationship>(
28        &mut self,
29        func: impl FnOnce(&mut RelatedSpawner<R>),
30    ) -> &mut Self {
31        let parent = self.id();
32        self.world_scope(|world| {
33            func(&mut RelatedSpawner::new(world, parent));
34        });
35        self
36    }
37
38    /// Relates the given entities to this entity with the relation `R`.
39    ///
40    /// See [`add_one_related`](Self::add_one_related) if you want relate only one entity.
41    pub fn add_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
42        let id = self.id();
43        self.world_scope(|world| {
44            for related in related {
45                world
46                    .entity_mut(*related)
47                    .modify_or_insert_relation_with_relationship_hook_mode::<R>(
48                        id,
49                        RelationshipHookMode::Run,
50                    );
51            }
52        });
53        self
54    }
55
56    /// Removes the relation `R` between this entity and all its related entities.
57    pub fn clear_related<R: Relationship>(&mut self) -> &mut Self {
58        self.remove::<R::RelationshipTarget>()
59    }
60
61    /// Relates the given entities to this entity with the relation `R`, starting at this particular index.
62    ///
63    /// If the `related` has duplicates, a related entity will take the index of its last occurrence in `related`.
64    /// If the indices go out of bounds, they will be clamped into bounds.
65    /// This will not re-order existing related entities unless they are in `related`.
66    ///
67    /// # Example
68    ///
69    /// ```
70    /// use bevy_ecs::prelude::*;
71    ///
72    /// let mut world = World::new();
73    /// let e0 = world.spawn_empty().id();
74    /// let e1 = world.spawn_empty().id();
75    /// let e2 = world.spawn_empty().id();
76    /// let e3 = world.spawn_empty().id();
77    /// let e4 = world.spawn_empty().id();
78    ///
79    /// let mut main_entity = world.spawn_empty();
80    /// main_entity.add_related::<ChildOf>(&[e0, e1, e2, e2]);
81    /// main_entity.insert_related::<ChildOf>(1, &[e0, e3, e4, e4]);
82    /// let main_id = main_entity.id();
83    ///
84    /// let relationship_source = main_entity.get::<Children>().unwrap().collection();
85    /// assert_eq!(relationship_source, &[e1, e0, e3, e2, e4]);
86    /// ```
87    pub fn insert_related<R: Relationship>(&mut self, index: usize, related: &[Entity]) -> &mut Self
88    where
89        <R::RelationshipTarget as RelationshipTarget>::Collection:
90            OrderedRelationshipSourceCollection,
91    {
92        let id = self.id();
93        self.world_scope(|world| {
94            for (offset, related) in related.iter().enumerate() {
95                let index = index.saturating_add(offset);
96                if world
97                    .get::<R>(*related)
98                    .is_some_and(|relationship| relationship.get() == id)
99                {
100                    world
101                        .get_mut::<R::RelationshipTarget>(id)
102                        .expect("hooks should have added relationship target")
103                        .collection_mut_risky()
104                        .place(*related, index);
105                } else {
106                    world
107                        .entity_mut(*related)
108                        .modify_or_insert_relation_with_relationship_hook_mode::<R>(
109                            id,
110                            RelationshipHookMode::Run,
111                        );
112                    world
113                        .get_mut::<R::RelationshipTarget>(id)
114                        .expect("hooks should have added relationship target")
115                        .collection_mut_risky()
116                        .place_most_recent(index);
117                }
118            }
119        });
120
121        self
122    }
123
124    /// Removes the relation `R` between this entity and the given entities.
125    pub fn remove_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
126        let id = self.id();
127        self.world_scope(|world| {
128            for related in related {
129                if world
130                    .get::<R>(*related)
131                    .is_some_and(|relationship| relationship.get() == id)
132                {
133                    world.entity_mut(*related).remove::<R>();
134                }
135            }
136        });
137
138        self
139    }
140
141    /// Replaces all the related entities with a new set of entities.
142    pub fn replace_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
143        type Collection<R> =
144            <<R as Relationship>::RelationshipTarget as RelationshipTarget>::Collection;
145
146        if related.is_empty() {
147            self.remove::<R::RelationshipTarget>();
148
149            return self;
150        }
151
152        let Some(existing_relations) = self.get_mut::<R::RelationshipTarget>() else {
153            return self.add_related::<R>(related);
154        };
155
156        // We replace the component here with a dummy value so we can modify it without taking it (this would create archetype move).
157        // SAFETY: We eventually return the correctly initialized collection into the target.
158        let mut relations = mem::replace(
159            existing_relations.into_inner(),
160            <R as Relationship>::RelationshipTarget::from_collection_risky(
161                Collection::<R>::with_capacity(0),
162            ),
163        );
164
165        let collection = relations.collection_mut_risky();
166
167        let mut potential_relations = EntityHashSet::from_iter(related.iter().copied());
168
169        let id = self.id();
170        self.world_scope(|world| {
171            for related in collection.iter() {
172                if !potential_relations.remove(related) {
173                    world.entity_mut(related).remove::<R>();
174                }
175            }
176
177            for related in potential_relations {
178                // SAFETY: We'll manually be adjusting the contents of the `RelationshipTarget` to fit the final state.
179                world
180                    .entity_mut(related)
181                    .modify_or_insert_relation_with_relationship_hook_mode::<R>(
182                        id,
183                        RelationshipHookMode::Skip,
184                    );
185            }
186        });
187
188        // SAFETY: The entities we're inserting will be the entities that were either already there or entities that we've just inserted.
189        collection.clear();
190        collection.extend_from_iter(related.iter().copied());
191        self.insert(relations);
192
193        self
194    }
195
196    /// Replaces all the related entities with a new set of entities.
197    ///
198    /// This is a more efficient of [`Self::replace_related`] which doesn't allocate.
199    /// The passed in arguments must adhere to these invariants:
200    /// - `entities_to_unrelate`: A slice of entities to remove from the relationship source.
201    ///   Entities need not be related to this entity, but must not appear in `entities_to_relate`
202    /// - `entities_to_relate`: A slice of entities to relate to this entity.
203    ///   This must contain all entities that will remain related (i.e. not those in `entities_to_unrelate`) plus the newly related entities.
204    /// - `newly_related_entities`: A subset of `entities_to_relate` containing only entities not already related to this entity.
205    /// - Slices **must not** contain any duplicates
206    ///
207    /// # Warning
208    ///
209    /// Violating these invariants may lead to panics, crashes or unpredictable engine behavior.
210    ///
211    /// # Panics
212    ///
213    /// Panics when debug assertions are enabled and any invariants are broken.
214    ///
215    // TODO: Consider making these iterators so users aren't required to allocate a separate buffers for the different slices.
216    pub fn replace_related_with_difference<R: Relationship>(
217        &mut self,
218        entities_to_unrelate: &[Entity],
219        entities_to_relate: &[Entity],
220        newly_related_entities: &[Entity],
221    ) -> &mut Self {
222        #[cfg(debug_assertions)]
223        {
224            let entities_to_relate = EntityHashSet::from_iter(entities_to_relate.iter().copied());
225            let entities_to_unrelate =
226                EntityHashSet::from_iter(entities_to_unrelate.iter().copied());
227            let mut newly_related_entities =
228                EntityHashSet::from_iter(newly_related_entities.iter().copied());
229            assert!(
230                entities_to_relate.is_disjoint(&entities_to_unrelate),
231                "`entities_to_relate` ({entities_to_relate:?}) shared entities with `entities_to_unrelate` ({entities_to_unrelate:?})"
232            );
233            assert!(
234                newly_related_entities.is_disjoint(&entities_to_unrelate),
235                "`newly_related_entities` ({newly_related_entities:?}) shared entities with `entities_to_unrelate ({entities_to_unrelate:?})`"
236            );
237            assert!(
238                newly_related_entities.is_subset(&entities_to_relate),
239                "`newly_related_entities` ({newly_related_entities:?}) wasn't a subset of `entities_to_relate` ({entities_to_relate:?})"
240            );
241
242            if let Some(target) = self.get::<R::RelationshipTarget>() {
243                let existing_relationships: EntityHashSet = target.collection().iter().collect();
244
245                assert!(
246                    existing_relationships.is_disjoint(&newly_related_entities),
247                    "`newly_related_entities` contains an entity that wouldn't be newly related"
248                );
249
250                newly_related_entities.extend(existing_relationships);
251                newly_related_entities -= &entities_to_unrelate;
252            }
253
254            assert_eq!(newly_related_entities, entities_to_relate, "`entities_to_relate` ({entities_to_relate:?}) didn't contain all entities that would end up related");
255        };
256
257        match self.get_mut::<R::RelationshipTarget>() {
258            None => {
259                self.add_related::<R>(entities_to_relate);
260
261                return self;
262            }
263            Some(mut target) => {
264                // SAFETY: The invariants expected by this function mean we'll only be inserting entities that are already related.
265                let collection = target.collection_mut_risky();
266                collection.clear();
267
268                collection.extend_from_iter(entities_to_relate.iter().copied());
269            }
270        }
271
272        let this = self.id();
273        self.world_scope(|world| {
274            for unrelate in entities_to_unrelate {
275                world.entity_mut(*unrelate).remove::<R>();
276            }
277
278            for new_relation in newly_related_entities {
279                // We changed the target collection manually so don't run the insert hook
280                world
281                    .entity_mut(*new_relation)
282                    .modify_or_insert_relation_with_relationship_hook_mode::<R>(
283                        this,
284                        RelationshipHookMode::Skip,
285                    );
286            }
287        });
288
289        self
290    }
291
292    /// Relates the given entity to this with the relation `R`.
293    ///
294    /// See [`add_related`](Self::add_related) if you want to relate more than one entity.
295    pub fn add_one_related<R: Relationship>(&mut self, entity: Entity) -> &mut Self {
296        self.add_related::<R>(&[entity])
297    }
298
299    /// Despawns entities that relate to this one via the given [`RelationshipTarget`].
300    /// This entity will not be despawned.
301    pub fn despawn_related<S: RelationshipTarget>(&mut self) -> &mut Self {
302        if let Some(sources) = self.get::<S>() {
303            // We have to collect here to defer removal, allowing observers and hooks to see this data
304            // before it is finally removed.
305            let sources = sources.iter().collect::<Vec<_>>();
306            self.world_scope(|world| {
307                for entity in sources {
308                    if let Ok(entity_mut) = world.get_entity_mut(entity) {
309                        entity_mut.despawn();
310                    };
311                }
312            });
313        }
314        self
315    }
316
317    /// Despawns the children of this entity.
318    /// This entity will not be despawned.
319    ///
320    /// This is a specialization of [`despawn_related`](EntityWorldMut::despawn_related), a more general method for despawning via relationships.
321    pub fn despawn_children(&mut self) -> &mut Self {
322        self.despawn_related::<Children>();
323        self
324    }
325
326    /// Inserts a component or bundle of components into the entity and all related entities,
327    /// traversing the relationship tracked in `S` in a breadth-first manner.
328    ///
329    /// # Warning
330    ///
331    /// This method should only be called on relationships that form a tree-like structure.
332    /// Any cycles will cause this method to loop infinitely.
333    // We could keep track of a list of visited entities and track cycles,
334    // but this is not a very well-defined operation (or hard to write) for arbitrary relationships.
335    pub fn insert_recursive<S: RelationshipTarget>(
336        &mut self,
337        bundle: impl Bundle + Clone,
338    ) -> &mut Self {
339        self.insert(bundle.clone());
340        if let Some(relationship_target) = self.get::<S>() {
341            let related_vec: Vec<Entity> = relationship_target.iter().collect();
342            for related in related_vec {
343                self.world_scope(|world| {
344                    world
345                        .entity_mut(related)
346                        .insert_recursive::<S>(bundle.clone());
347                });
348            }
349        }
350
351        self
352    }
353
354    /// Removes a component or bundle of components of type `B` from the entity and all related entities,
355    /// traversing the relationship tracked in `S` in a breadth-first manner.
356    ///
357    /// # Warning
358    ///
359    /// This method should only be called on relationships that form a tree-like structure.
360    /// Any cycles will cause this method to loop infinitely.
361    pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
362        self.remove::<B>();
363        if let Some(relationship_target) = self.get::<S>() {
364            let related_vec: Vec<Entity> = relationship_target.iter().collect();
365            for related in related_vec {
366                self.world_scope(|world| {
367                    world.entity_mut(related).remove_recursive::<S, B>();
368                });
369            }
370        }
371
372        self
373    }
374
375    fn modify_or_insert_relation_with_relationship_hook_mode<R: Relationship>(
376        &mut self,
377        entity: Entity,
378        relationship_hook_mode: RelationshipHookMode,
379    ) {
380        // Check if the relation edge holds additional data
381        if size_of::<R>() > size_of::<Entity>() {
382            self.assert_not_despawned();
383
384            let this = self.id();
385
386            let modified = self.world_scope(|world| {
387                let modified = DeferredWorld::from(&mut *world)
388                    .modify_component_with_relationship_hook_mode::<R, _>(
389                        this,
390                        relationship_hook_mode,
391                        |r| r.set_risky(entity),
392                    )
393                    .expect("entity access must be valid")
394                    .is_some();
395
396                world.flush();
397
398                modified
399            });
400
401            if modified {
402                return;
403            }
404        }
405
406        self.insert_with_relationship_hook_mode(R::from(entity), relationship_hook_mode);
407    }
408}
409
410impl<'a> EntityCommands<'a> {
411    /// Spawns a entity related to this entity (with the `R` relationship) by taking a bundle
412    pub fn with_related<R: Relationship>(&mut self, bundle: impl Bundle) -> &mut Self {
413        let parent = self.id();
414        self.commands.spawn((bundle, R::from(parent)));
415        self
416    }
417
418    /// Spawns entities related to this entity (with the `R` relationship) by taking a function that operates on a [`RelatedSpawner`].
419    pub fn with_related_entities<R: Relationship>(
420        &mut self,
421        func: impl FnOnce(&mut RelatedSpawnerCommands<R>),
422    ) -> &mut Self {
423        let id = self.id();
424        func(&mut RelatedSpawnerCommands::new(self.commands(), id));
425        self
426    }
427
428    /// Relates the given entities to this entity with the relation `R`.
429    ///
430    /// See [`add_one_related`](Self::add_one_related) if you want relate only one entity.
431    pub fn add_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
432        let related: Box<[Entity]> = related.into();
433
434        self.queue(move |mut entity: EntityWorldMut| {
435            entity.add_related::<R>(&related);
436        })
437    }
438
439    /// Removes the relation `R` between this entity and all its related entities.
440    pub fn clear_related<R: Relationship>(&mut self) -> &mut Self {
441        self.queue(|mut entity: EntityWorldMut| {
442            entity.clear_related::<R>();
443        })
444    }
445
446    /// Relates the given entities to this entity with the relation `R`, starting at this particular index.
447    ///
448    /// If the `related` has duplicates, a related entity will take the index of its last occurrence in `related`.
449    /// If the indices go out of bounds, they will be clamped into bounds.
450    /// This will not re-order existing related entities unless they are in `related`.
451    pub fn insert_related<R: Relationship>(&mut self, index: usize, related: &[Entity]) -> &mut Self
452    where
453        <R::RelationshipTarget as RelationshipTarget>::Collection:
454            OrderedRelationshipSourceCollection,
455    {
456        let related: Box<[Entity]> = related.into();
457
458        self.queue(move |mut entity: EntityWorldMut| {
459            entity.insert_related::<R>(index, &related);
460        })
461    }
462
463    /// Relates the given entity to this with the relation `R`.
464    ///
465    /// See [`add_related`](Self::add_related) if you want to relate more than one entity.
466    pub fn add_one_related<R: Relationship>(&mut self, entity: Entity) -> &mut Self {
467        self.add_related::<R>(&[entity])
468    }
469
470    /// Removes the relation `R` between this entity and the given entities.
471    pub fn remove_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
472        let related: Box<[Entity]> = related.into();
473
474        self.queue(move |mut entity: EntityWorldMut| {
475            entity.remove_related::<R>(&related);
476        })
477    }
478
479    /// Replaces all the related entities with the given set of new related entities.
480    pub fn replace_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
481        let related: Box<[Entity]> = related.into();
482
483        self.queue(move |mut entity: EntityWorldMut| {
484            entity.replace_related::<R>(&related);
485        })
486    }
487
488    /// Replaces all the related entities with a new set of entities.
489    ///
490    /// # Warning
491    ///
492    /// Failing to maintain the functions invariants may lead to erratic engine behavior including random crashes.
493    /// Refer to [`EntityWorldMut::replace_related_with_difference`] for a list of these invariants.
494    ///
495    /// # Panics
496    ///
497    /// Panics when debug assertions are enable, an invariant is are broken and the command is executed.
498    pub fn replace_related_with_difference<R: Relationship>(
499        &mut self,
500        entities_to_unrelate: &[Entity],
501        entities_to_relate: &[Entity],
502        newly_related_entities: &[Entity],
503    ) -> &mut Self {
504        let entities_to_unrelate: Box<[Entity]> = entities_to_unrelate.into();
505        let entities_to_relate: Box<[Entity]> = entities_to_relate.into();
506        let newly_related_entities: Box<[Entity]> = newly_related_entities.into();
507
508        self.queue(move |mut entity: EntityWorldMut| {
509            entity.replace_related_with_difference::<R>(
510                &entities_to_unrelate,
511                &entities_to_relate,
512                &newly_related_entities,
513            );
514        })
515    }
516
517    /// Despawns entities that relate to this one via the given [`RelationshipTarget`].
518    /// This entity will not be despawned.
519    pub fn despawn_related<S: RelationshipTarget>(&mut self) -> &mut Self {
520        self.queue(move |mut entity: EntityWorldMut| {
521            entity.despawn_related::<S>();
522        })
523    }
524
525    /// Despawns the children of this entity.
526    /// This entity will not be despawned.
527    ///
528    /// This is a specialization of [`despawn_related`](EntityCommands::despawn_related), a more general method for despawning via relationships.
529    pub fn despawn_children(&mut self) -> &mut Self {
530        self.despawn_related::<Children>()
531    }
532
533    /// Inserts a component or bundle of components into the entity and all related entities,
534    /// traversing the relationship tracked in `S` in a breadth-first manner.
535    ///
536    /// # Warning
537    ///
538    /// This method should only be called on relationships that form a tree-like structure.
539    /// Any cycles will cause this method to loop infinitely.
540    pub fn insert_recursive<S: RelationshipTarget>(
541        &mut self,
542        bundle: impl Bundle + Clone,
543    ) -> &mut Self {
544        self.queue(move |mut entity: EntityWorldMut| {
545            entity.insert_recursive::<S>(bundle);
546        })
547    }
548
549    /// Removes a component or bundle of components of type `B` from the entity and all related entities,
550    /// traversing the relationship tracked in `S` in a breadth-first manner.
551    ///
552    /// # Warning
553    ///
554    /// This method should only be called on relationships that form a tree-like structure.
555    /// Any cycles will cause this method to loop infinitely.
556    pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
557        self.queue(move |mut entity: EntityWorldMut| {
558            entity.remove_recursive::<S, B>();
559        })
560    }
561}
562
563/// Directly spawns related "source" entities with the given [`Relationship`], targeting
564/// a specific entity.
565pub struct RelatedSpawner<'w, R: Relationship> {
566    target: Entity,
567    world: &'w mut World,
568    _marker: PhantomData<R>,
569}
570
571impl<'w, R: Relationship> RelatedSpawner<'w, R> {
572    /// Creates a new instance that will spawn entities targeting the `target` entity.
573    pub fn new(world: &'w mut World, target: Entity) -> Self {
574        Self {
575            world,
576            target,
577            _marker: PhantomData,
578        }
579    }
580
581    /// Spawns an entity with the given `bundle` and an `R` relationship targeting the `target`
582    /// entity this spawner was initialized with.
583    pub fn spawn(&mut self, bundle: impl Bundle) -> EntityWorldMut<'_> {
584        self.world.spawn((R::from(self.target), bundle))
585    }
586
587    /// Spawns an entity with an `R` relationship targeting the `target`
588    /// entity this spawner was initialized with.
589    pub fn spawn_empty(&mut self) -> EntityWorldMut<'_> {
590        self.world.spawn(R::from(self.target))
591    }
592
593    /// Returns the "target entity" used when spawning entities with an `R` [`Relationship`].
594    pub fn target_entity(&self) -> Entity {
595        self.target
596    }
597
598    /// Returns a reference to the underlying [`World`].
599    pub fn world(&self) -> &World {
600        self.world
601    }
602
603    /// Returns a mutable reference to the underlying [`World`].
604    pub fn world_mut(&mut self) -> &mut World {
605        self.world
606    }
607}
608
609/// Uses commands to spawn related "source" entities with the given [`Relationship`], targeting
610/// a specific entity.
611pub struct RelatedSpawnerCommands<'w, R: Relationship> {
612    target: Entity,
613    commands: Commands<'w, 'w>,
614    _marker: PhantomData<R>,
615}
616
617impl<'w, R: Relationship> RelatedSpawnerCommands<'w, R> {
618    /// Creates a new instance that will spawn entities targeting the `target` entity.
619    pub fn new(commands: Commands<'w, 'w>, target: Entity) -> Self {
620        Self {
621            commands,
622            target,
623            _marker: PhantomData,
624        }
625    }
626
627    /// Spawns an entity with the given `bundle` and an `R` relationship targeting the `target`
628    /// entity this spawner was initialized with.
629    pub fn spawn(&mut self, bundle: impl Bundle) -> EntityCommands<'_> {
630        self.commands.spawn((R::from(self.target), bundle))
631    }
632
633    /// Spawns an entity with an `R` relationship targeting the `target`
634    /// entity this spawner was initialized with.
635    pub fn spawn_empty(&mut self) -> EntityCommands<'_> {
636        self.commands.spawn(R::from(self.target))
637    }
638
639    /// Returns the "target entity" used when spawning entities with an `R` [`Relationship`].
640    pub fn target_entity(&self) -> Entity {
641        self.target
642    }
643
644    /// Returns the underlying [`Commands`].
645    pub fn commands(&mut self) -> Commands<'_, '_> {
646        self.commands.reborrow()
647    }
648
649    /// Returns a mutable reference to the underlying [`Commands`].
650    pub fn commands_mut(&mut self) -> &mut Commands<'w, 'w> {
651        &mut self.commands
652    }
653}
654
655#[cfg(test)]
656mod tests {
657    use super::*;
658    use crate::prelude::{ChildOf, Children, Component};
659
660    #[derive(Component, Clone, Copy)]
661    struct TestComponent;
662
663    #[test]
664    fn insert_and_remove_recursive() {
665        let mut world = World::new();
666
667        let a = world.spawn_empty().id();
668        let b = world.spawn(ChildOf(a)).id();
669        let c = world.spawn(ChildOf(a)).id();
670        let d = world.spawn(ChildOf(b)).id();
671
672        world
673            .entity_mut(a)
674            .insert_recursive::<Children>(TestComponent);
675
676        for entity in [a, b, c, d] {
677            assert!(world.entity(entity).contains::<TestComponent>());
678        }
679
680        world
681            .entity_mut(b)
682            .remove_recursive::<Children, TestComponent>();
683
684        // Parent
685        assert!(world.entity(a).contains::<TestComponent>());
686        // Target
687        assert!(!world.entity(b).contains::<TestComponent>());
688        // Sibling
689        assert!(world.entity(c).contains::<TestComponent>());
690        // Child
691        assert!(!world.entity(d).contains::<TestComponent>());
692
693        world
694            .entity_mut(a)
695            .remove_recursive::<Children, TestComponent>();
696
697        for entity in [a, b, c, d] {
698            assert!(!world.entity(entity).contains::<TestComponent>());
699        }
700    }
701
702    #[test]
703    fn remove_all_related() {
704        let mut world = World::new();
705
706        let a = world.spawn_empty().id();
707        let b = world.spawn(ChildOf(a)).id();
708        let c = world.spawn(ChildOf(a)).id();
709
710        world.entity_mut(a).clear_related::<ChildOf>();
711
712        assert_eq!(world.entity(a).get::<Children>(), None);
713        assert_eq!(world.entity(b).get::<ChildOf>(), None);
714        assert_eq!(world.entity(c).get::<ChildOf>(), None);
715    }
716
717    #[test]
718    fn replace_related_works() {
719        let mut world = World::new();
720        let child1 = world.spawn_empty().id();
721        let child2 = world.spawn_empty().id();
722        let child3 = world.spawn_empty().id();
723
724        let mut parent = world.spawn_empty();
725        parent.add_children(&[child1, child2]);
726        let child_value = ChildOf(parent.id());
727        let some_child = Some(&child_value);
728
729        parent.replace_children(&[child2, child3]);
730        let children = parent.get::<Children>().unwrap().collection();
731        assert_eq!(children, &[child2, child3]);
732        assert_eq!(parent.world().get::<ChildOf>(child1), None);
733        assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
734        assert_eq!(parent.world().get::<ChildOf>(child3), some_child);
735
736        parent.replace_children_with_difference(&[child3], &[child1, child2], &[child1]);
737        let children = parent.get::<Children>().unwrap().collection();
738        assert_eq!(children, &[child1, child2]);
739        assert_eq!(parent.world().get::<ChildOf>(child1), some_child);
740        assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
741        assert_eq!(parent.world().get::<ChildOf>(child3), None);
742    }
743
744    #[test]
745    fn add_related_keeps_relationship_data() {
746        #[derive(Component, PartialEq, Debug)]
747        #[relationship(relationship_target = Parent)]
748        struct Child {
749            #[relationship]
750            parent: Entity,
751            data: u8,
752        }
753
754        #[derive(Component)]
755        #[relationship_target(relationship = Child)]
756        struct Parent(Vec<Entity>);
757
758        let mut world = World::new();
759        let parent1 = world.spawn_empty().id();
760        let parent2 = world.spawn_empty().id();
761        let child = world
762            .spawn(Child {
763                parent: parent1,
764                data: 42,
765            })
766            .id();
767
768        world.entity_mut(parent2).add_related::<Child>(&[child]);
769        assert_eq!(
770            world.get::<Child>(child),
771            Some(&Child {
772                parent: parent2,
773                data: 42
774            })
775        );
776    }
777
778    #[test]
779    fn insert_related_keeps_relationship_data() {
780        #[derive(Component, PartialEq, Debug)]
781        #[relationship(relationship_target = Parent)]
782        struct Child {
783            #[relationship]
784            parent: Entity,
785            data: u8,
786        }
787
788        #[derive(Component)]
789        #[relationship_target(relationship = Child)]
790        struct Parent(Vec<Entity>);
791
792        let mut world = World::new();
793        let parent1 = world.spawn_empty().id();
794        let parent2 = world.spawn_empty().id();
795        let child = world
796            .spawn(Child {
797                parent: parent1,
798                data: 42,
799            })
800            .id();
801
802        world
803            .entity_mut(parent2)
804            .insert_related::<Child>(0, &[child]);
805        assert_eq!(
806            world.get::<Child>(child),
807            Some(&Child {
808                parent: parent2,
809                data: 42
810            })
811        );
812    }
813
814    #[test]
815    fn replace_related_keeps_relationship_data() {
816        #[derive(Component, PartialEq, Debug)]
817        #[relationship(relationship_target = Parent)]
818        struct Child {
819            #[relationship]
820            parent: Entity,
821            data: u8,
822        }
823
824        #[derive(Component)]
825        #[relationship_target(relationship = Child)]
826        struct Parent(Vec<Entity>);
827
828        let mut world = World::new();
829        let parent1 = world.spawn_empty().id();
830        let parent2 = world.spawn_empty().id();
831        let child = world
832            .spawn(Child {
833                parent: parent1,
834                data: 42,
835            })
836            .id();
837
838        world
839            .entity_mut(parent2)
840            .replace_related_with_difference::<Child>(&[], &[child], &[child]);
841        assert_eq!(
842            world.get::<Child>(child),
843            Some(&Child {
844                parent: parent2,
845                data: 42
846            })
847        );
848
849        world.entity_mut(parent1).replace_related::<Child>(&[child]);
850        assert_eq!(
851            world.get::<Child>(child),
852            Some(&Child {
853                parent: parent1,
854                data: 42
855            })
856        );
857    }
858
859    #[test]
860    fn replace_related_keeps_relationship_target_data() {
861        #[derive(Component)]
862        #[relationship(relationship_target = Parent)]
863        struct Child(Entity);
864
865        #[derive(Component)]
866        #[relationship_target(relationship = Child)]
867        struct Parent {
868            #[relationship]
869            children: Vec<Entity>,
870            data: u8,
871        }
872
873        let mut world = World::new();
874        let child1 = world.spawn_empty().id();
875        let child2 = world.spawn_empty().id();
876        let mut parent = world.spawn_empty();
877        parent.add_related::<Child>(&[child1]);
878        parent.get_mut::<Parent>().unwrap().data = 42;
879
880        parent.replace_related_with_difference::<Child>(&[child1], &[child2], &[child2]);
881        let data = parent.get::<Parent>().unwrap().data;
882        assert_eq!(data, 42);
883
884        parent.replace_related::<Child>(&[child1]);
885        let data = parent.get::<Parent>().unwrap().data;
886        assert_eq!(data, 42);
887    }
888
889    #[test]
890    fn despawn_related_observers_can_access_relationship_data() {
891        use crate::lifecycle::Replace;
892        use crate::observer::On;
893        use crate::prelude::Has;
894        use crate::system::Query;
895
896        #[derive(Component)]
897        struct MyComponent;
898
899        #[derive(Component, Default)]
900        struct ObserverResult {
901            success: bool,
902        }
903
904        let mut world = World::new();
905        let result_entity = world.spawn(ObserverResult::default()).id();
906
907        world.add_observer(
908            move |replace: On<Replace, MyComponent>,
909                  has_relationship: Query<Has<ChildOf>>,
910                  mut results: Query<&mut ObserverResult>| {
911                if has_relationship.get(replace.entity).unwrap_or(false) {
912                    results.get_mut(result_entity).unwrap().success = true;
913                }
914            },
915        );
916
917        let parent = world.spawn_empty().id();
918        let _child = world.spawn((MyComponent, ChildOf(parent))).id();
919
920        world.entity_mut(parent).despawn_related::<Children>();
921
922        assert!(world.get::<ObserverResult>(result_entity).unwrap().success);
923    }
924}