Skip to main content

bevy_ecs/relationship/
related_methods.rs

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