bevy_ecs/relationship/
relationship_source_collection.rs

1use alloc::collections::{btree_set, BTreeSet};
2use core::{
3    hash::BuildHasher,
4    ops::{Deref, DerefMut},
5};
6
7use crate::entity::{Entity, EntityHashSet, EntityIndexSet};
8use alloc::vec::Vec;
9use indexmap::IndexSet;
10use smallvec::SmallVec;
11
12/// The internal [`Entity`] collection used by a [`RelationshipTarget`](crate::relationship::RelationshipTarget) component.
13/// This is not intended to be modified directly by users, as it could invalidate the correctness of relationships.
14pub trait RelationshipSourceCollection {
15    /// The type of iterator returned by the `iter` method.
16    ///
17    /// This is an associated type (rather than using a method that returns an opaque return-position impl trait)
18    /// to ensure that all methods and traits (like [`DoubleEndedIterator`]) of the underlying collection's iterator
19    /// are available to the user when implemented without unduly restricting the possible collections.
20    ///
21    /// The [`SourceIter`](super::SourceIter) type alias can be helpful to reduce confusion when working with this associated type.
22    type SourceIter<'a>: Iterator<Item = Entity>
23    where
24        Self: 'a;
25
26    /// Creates a new empty instance.
27    fn new() -> Self;
28
29    /// Returns an instance with the given pre-allocated entity `capacity`.
30    ///
31    /// Some collections will ignore the provided `capacity` and return a default instance.
32    fn with_capacity(capacity: usize) -> Self;
33
34    /// Reserves capacity for at least `additional` more entities to be inserted.
35    ///
36    /// Not all collections support this operation, in which case it is a no-op.
37    fn reserve(&mut self, additional: usize);
38
39    /// Adds the given `entity` to the collection.
40    ///
41    /// Returns whether the entity was added to the collection.
42    /// Mainly useful when dealing with collections that don't allow
43    /// multiple instances of the same entity ([`EntityHashSet`]).
44    fn add(&mut self, entity: Entity) -> bool;
45
46    /// Removes the given `entity` from the collection.
47    ///
48    /// Returns whether the collection actually contained
49    /// the entity.
50    fn remove(&mut self, entity: Entity) -> bool;
51
52    /// Iterates all entities in the collection.
53    fn iter(&self) -> Self::SourceIter<'_>;
54
55    /// Returns the current length of the collection.
56    fn len(&self) -> usize;
57
58    /// Clears the collection.
59    fn clear(&mut self);
60
61    /// Attempts to save memory by shrinking the capacity to fit the current length.
62    ///
63    /// This operation is a no-op for collections that do not support it.
64    fn shrink_to_fit(&mut self);
65
66    /// Returns true if the collection contains no entities.
67    #[inline]
68    fn is_empty(&self) -> bool {
69        self.len() == 0
70    }
71
72    /// For one-to-one relationships, returns the entity that should be removed before adding a new one.
73    /// Returns `None` for one-to-many relationships or when no entity needs to be removed.
74    fn source_to_remove_before_add(&self) -> Option<Entity> {
75        None
76    }
77
78    /// Add multiple entities to collection at once.
79    ///
80    /// May be faster than repeatedly calling [`Self::add`].
81    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>);
82}
83
84/// This trait signals that a [`RelationshipSourceCollection`] is ordered.
85pub trait OrderedRelationshipSourceCollection: RelationshipSourceCollection {
86    /// Inserts the entity at a specific index.
87    /// If the index is too large, the entity will be added to the end of the collection.
88    fn insert(&mut self, index: usize, entity: Entity);
89    /// Removes the entity at the specified index if it exists.
90    fn remove_at(&mut self, index: usize) -> Option<Entity>;
91    /// Inserts the entity at a specific index.
92    /// This will never reorder other entities.
93    /// If the index is too large, the entity will be added to the end of the collection.
94    fn insert_stable(&mut self, index: usize, entity: Entity);
95    /// Removes the entity at the specified index if it exists.
96    /// This will never reorder other entities.
97    fn remove_at_stable(&mut self, index: usize) -> Option<Entity>;
98    /// Sorts the source collection.
99    fn sort(&mut self);
100    /// Inserts the entity at the proper place to maintain sorting.
101    fn insert_sorted(&mut self, entity: Entity);
102
103    /// This places the most recently added entity at the particular index.
104    fn place_most_recent(&mut self, index: usize);
105
106    /// This places the given entity at the particular index.
107    /// This will do nothing if the entity is not in the collection.
108    /// If the index is out of bounds, this will put the entity at the end.
109    fn place(&mut self, entity: Entity, index: usize);
110
111    /// Adds the entity at index 0.
112    fn push_front(&mut self, entity: Entity) {
113        self.insert(0, entity);
114    }
115
116    /// Adds the entity to the back of the collection.
117    fn push_back(&mut self, entity: Entity) {
118        self.insert(usize::MAX, entity);
119    }
120
121    /// Removes the first entity.
122    fn pop_front(&mut self) -> Option<Entity> {
123        self.remove_at(0)
124    }
125
126    /// Removes the last entity.
127    fn pop_back(&mut self) -> Option<Entity> {
128        if self.is_empty() {
129            None
130        } else {
131            self.remove_at(self.len() - 1)
132        }
133    }
134}
135
136impl RelationshipSourceCollection for Vec<Entity> {
137    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
138
139    fn new() -> Self {
140        Vec::new()
141    }
142
143    fn reserve(&mut self, additional: usize) {
144        Vec::reserve(self, additional);
145    }
146
147    fn with_capacity(capacity: usize) -> Self {
148        Vec::with_capacity(capacity)
149    }
150
151    fn add(&mut self, entity: Entity) -> bool {
152        Vec::push(self, entity);
153
154        true
155    }
156
157    fn remove(&mut self, entity: Entity) -> bool {
158        if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
159            Vec::remove(self, index);
160            return true;
161        }
162
163        false
164    }
165
166    fn iter(&self) -> Self::SourceIter<'_> {
167        <[Entity]>::iter(self).copied()
168    }
169
170    fn len(&self) -> usize {
171        Vec::len(self)
172    }
173
174    fn clear(&mut self) {
175        self.clear();
176    }
177
178    fn shrink_to_fit(&mut self) {
179        Vec::shrink_to_fit(self);
180    }
181
182    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
183        self.extend(entities);
184    }
185}
186
187impl OrderedRelationshipSourceCollection for Vec<Entity> {
188    fn insert(&mut self, index: usize, entity: Entity) {
189        self.push(entity);
190        let len = self.len();
191        if index < len {
192            self.swap(index, len - 1);
193        }
194    }
195
196    fn remove_at(&mut self, index: usize) -> Option<Entity> {
197        (index < self.len()).then(|| self.swap_remove(index))
198    }
199
200    fn insert_stable(&mut self, index: usize, entity: Entity) {
201        if index < self.len() {
202            Vec::insert(self, index, entity);
203        } else {
204            self.push(entity);
205        }
206    }
207
208    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
209        (index < self.len()).then(|| self.remove(index))
210    }
211
212    fn sort(&mut self) {
213        self.sort_unstable();
214    }
215
216    fn insert_sorted(&mut self, entity: Entity) {
217        let index = self.partition_point(|e| e <= &entity);
218        self.insert_stable(index, entity);
219    }
220
221    fn place_most_recent(&mut self, index: usize) {
222        if let Some(entity) = self.pop() {
223            let index = index.min(self.len());
224            self.insert(index, entity);
225        }
226    }
227
228    fn place(&mut self, entity: Entity, index: usize) {
229        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
230            let index = index.min(self.len());
231            Vec::remove(self, current);
232            self.insert(index, entity);
233        };
234    }
235}
236
237impl RelationshipSourceCollection for EntityHashSet {
238    type SourceIter<'a> = core::iter::Copied<crate::entity::hash_set::Iter<'a>>;
239
240    fn new() -> Self {
241        EntityHashSet::new()
242    }
243
244    fn reserve(&mut self, additional: usize) {
245        self.0.reserve(additional);
246    }
247
248    fn with_capacity(capacity: usize) -> Self {
249        EntityHashSet::with_capacity(capacity)
250    }
251
252    fn add(&mut self, entity: Entity) -> bool {
253        self.insert(entity)
254    }
255
256    fn remove(&mut self, entity: Entity) -> bool {
257        // We need to call the remove method on the underlying hash set,
258        // which takes its argument by reference
259        self.0.remove(&entity)
260    }
261
262    fn iter(&self) -> Self::SourceIter<'_> {
263        self.iter().copied()
264    }
265
266    fn len(&self) -> usize {
267        self.len()
268    }
269
270    fn clear(&mut self) {
271        self.0.clear();
272    }
273
274    fn shrink_to_fit(&mut self) {
275        self.0.shrink_to_fit();
276    }
277
278    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
279        self.extend(entities);
280    }
281}
282
283impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
284    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
285
286    fn new() -> Self {
287        SmallVec::new()
288    }
289
290    fn reserve(&mut self, additional: usize) {
291        SmallVec::reserve(self, additional);
292    }
293
294    fn with_capacity(capacity: usize) -> Self {
295        SmallVec::with_capacity(capacity)
296    }
297
298    fn add(&mut self, entity: Entity) -> bool {
299        SmallVec::push(self, entity);
300
301        true
302    }
303
304    fn remove(&mut self, entity: Entity) -> bool {
305        if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
306            SmallVec::remove(self, index);
307            return true;
308        }
309
310        false
311    }
312
313    fn iter(&self) -> Self::SourceIter<'_> {
314        <[Entity]>::iter(self).copied()
315    }
316
317    fn len(&self) -> usize {
318        SmallVec::len(self)
319    }
320
321    fn clear(&mut self) {
322        self.clear();
323    }
324
325    fn shrink_to_fit(&mut self) {
326        SmallVec::shrink_to_fit(self);
327    }
328
329    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
330        self.extend(entities);
331    }
332}
333
334impl RelationshipSourceCollection for Entity {
335    type SourceIter<'a> = core::option::IntoIter<Entity>;
336
337    fn new() -> Self {
338        Entity::PLACEHOLDER
339    }
340
341    fn reserve(&mut self, _: usize) {}
342
343    fn with_capacity(_capacity: usize) -> Self {
344        Self::new()
345    }
346
347    fn add(&mut self, entity: Entity) -> bool {
348        *self = entity;
349        true
350    }
351
352    fn remove(&mut self, entity: Entity) -> bool {
353        if *self == entity {
354            *self = Entity::PLACEHOLDER;
355
356            return true;
357        }
358
359        false
360    }
361
362    fn iter(&self) -> Self::SourceIter<'_> {
363        if *self == Entity::PLACEHOLDER {
364            None.into_iter()
365        } else {
366            Some(*self).into_iter()
367        }
368    }
369
370    fn len(&self) -> usize {
371        if *self == Entity::PLACEHOLDER {
372            return 0;
373        }
374        1
375    }
376
377    fn clear(&mut self) {
378        *self = Entity::PLACEHOLDER;
379    }
380
381    fn shrink_to_fit(&mut self) {}
382
383    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
384        for entity in entities {
385            *self = entity;
386        }
387    }
388
389    fn source_to_remove_before_add(&self) -> Option<Entity> {
390        if *self != Entity::PLACEHOLDER {
391            Some(*self)
392        } else {
393            None
394        }
395    }
396}
397
398impl<const N: usize> OrderedRelationshipSourceCollection for SmallVec<[Entity; N]> {
399    fn insert(&mut self, index: usize, entity: Entity) {
400        self.push(entity);
401        let len = self.len();
402        if index < len {
403            self.swap(index, len - 1);
404        }
405    }
406
407    fn remove_at(&mut self, index: usize) -> Option<Entity> {
408        (index < self.len()).then(|| self.swap_remove(index))
409    }
410
411    fn insert_stable(&mut self, index: usize, entity: Entity) {
412        if index < self.len() {
413            SmallVec::<[Entity; N]>::insert(self, index, entity);
414        } else {
415            self.push(entity);
416        }
417    }
418
419    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
420        (index < self.len()).then(|| self.remove(index))
421    }
422
423    fn sort(&mut self) {
424        self.sort_unstable();
425    }
426
427    fn insert_sorted(&mut self, entity: Entity) {
428        let index = self.partition_point(|e| e <= &entity);
429        self.insert_stable(index, entity);
430    }
431
432    fn place_most_recent(&mut self, index: usize) {
433        if let Some(entity) = self.pop() {
434            let index = index.min(self.len() - 1);
435            self.insert(index, entity);
436        }
437    }
438
439    fn place(&mut self, entity: Entity, index: usize) {
440        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
441            // The len is at least 1, so the subtraction is safe.
442            let index = index.min(self.len() - 1);
443            SmallVec::<[Entity; N]>::remove(self, current);
444            self.insert(index, entity);
445        };
446    }
447}
448
449impl<S: BuildHasher + Default> RelationshipSourceCollection for IndexSet<Entity, S> {
450    type SourceIter<'a>
451        = core::iter::Copied<indexmap::set::Iter<'a, Entity>>
452    where
453        S: 'a;
454
455    fn new() -> Self {
456        IndexSet::default()
457    }
458
459    fn reserve(&mut self, additional: usize) {
460        self.reserve(additional);
461    }
462
463    fn with_capacity(capacity: usize) -> Self {
464        IndexSet::with_capacity_and_hasher(capacity, S::default())
465    }
466
467    fn add(&mut self, entity: Entity) -> bool {
468        self.insert(entity)
469    }
470
471    fn remove(&mut self, entity: Entity) -> bool {
472        self.shift_remove(&entity)
473    }
474
475    fn iter(&self) -> Self::SourceIter<'_> {
476        self.iter().copied()
477    }
478
479    fn len(&self) -> usize {
480        self.len()
481    }
482
483    fn clear(&mut self) {
484        self.clear();
485    }
486
487    fn shrink_to_fit(&mut self) {
488        self.shrink_to_fit();
489    }
490
491    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
492        self.extend(entities);
493    }
494}
495
496impl RelationshipSourceCollection for EntityIndexSet {
497    type SourceIter<'a> = core::iter::Copied<crate::entity::index_set::Iter<'a>>;
498
499    fn new() -> Self {
500        EntityIndexSet::new()
501    }
502
503    fn reserve(&mut self, additional: usize) {
504        self.deref_mut().reserve(additional);
505    }
506
507    fn with_capacity(capacity: usize) -> Self {
508        EntityIndexSet::with_capacity(capacity)
509    }
510
511    fn add(&mut self, entity: Entity) -> bool {
512        self.insert(entity)
513    }
514
515    fn remove(&mut self, entity: Entity) -> bool {
516        self.deref_mut().shift_remove(&entity)
517    }
518
519    fn iter(&self) -> Self::SourceIter<'_> {
520        self.iter().copied()
521    }
522
523    fn len(&self) -> usize {
524        self.deref().len()
525    }
526
527    fn clear(&mut self) {
528        self.deref_mut().clear();
529    }
530
531    fn shrink_to_fit(&mut self) {
532        self.deref_mut().shrink_to_fit();
533    }
534
535    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
536        self.extend(entities);
537    }
538}
539
540impl RelationshipSourceCollection for BTreeSet<Entity> {
541    type SourceIter<'a> = core::iter::Copied<btree_set::Iter<'a, Entity>>;
542
543    fn new() -> Self {
544        BTreeSet::new()
545    }
546
547    fn with_capacity(_: usize) -> Self {
548        // BTreeSet doesn't have a capacity
549        Self::new()
550    }
551
552    fn reserve(&mut self, _: usize) {
553        // BTreeSet doesn't have a capacity
554    }
555
556    fn add(&mut self, entity: Entity) -> bool {
557        self.insert(entity)
558    }
559
560    fn remove(&mut self, entity: Entity) -> bool {
561        self.remove(&entity)
562    }
563
564    fn iter(&self) -> Self::SourceIter<'_> {
565        self.iter().copied()
566    }
567
568    fn len(&self) -> usize {
569        self.len()
570    }
571
572    fn clear(&mut self) {
573        self.clear();
574    }
575
576    fn shrink_to_fit(&mut self) {
577        // BTreeSet doesn't have a capacity
578    }
579
580    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
581        self.extend(entities);
582    }
583}
584
585#[cfg(test)]
586mod tests {
587    use super::*;
588    use crate::prelude::{Component, World};
589    use crate::relationship::RelationshipTarget;
590
591    #[test]
592    fn vec_relationship_source_collection() {
593        #[derive(Component)]
594        #[relationship(relationship_target = RelTarget)]
595        struct Rel(Entity);
596
597        #[derive(Component)]
598        #[relationship_target(relationship = Rel, linked_spawn)]
599        struct RelTarget(Vec<Entity>);
600
601        let mut world = World::new();
602        let a = world.spawn_empty().id();
603        let b = world.spawn_empty().id();
604
605        world.entity_mut(a).insert(Rel(b));
606
607        let rel_target = world.get::<RelTarget>(b).unwrap();
608        let collection = rel_target.collection();
609        assert_eq!(collection, &alloc::vec!(a));
610    }
611
612    #[test]
613    fn smallvec_relationship_source_collection() {
614        #[derive(Component)]
615        #[relationship(relationship_target = RelTarget)]
616        struct Rel(Entity);
617
618        #[derive(Component)]
619        #[relationship_target(relationship = Rel, linked_spawn)]
620        struct RelTarget(SmallVec<[Entity; 4]>);
621
622        let mut world = World::new();
623        let a = world.spawn_empty().id();
624        let b = world.spawn_empty().id();
625
626        world.entity_mut(a).insert(Rel(b));
627
628        let rel_target = world.get::<RelTarget>(b).unwrap();
629        let collection = rel_target.collection();
630        assert_eq!(collection, &SmallVec::from_buf([a]));
631    }
632
633    #[test]
634    fn entity_relationship_source_collection() {
635        #[derive(Component)]
636        #[relationship(relationship_target = RelTarget)]
637        struct Rel(Entity);
638
639        #[derive(Component)]
640        #[relationship_target(relationship = Rel)]
641        struct RelTarget(Entity);
642
643        let mut world = World::new();
644        let a = world.spawn_empty().id();
645        let b = world.spawn_empty().id();
646
647        world.entity_mut(a).insert(Rel(b));
648
649        let rel_target = world.get::<RelTarget>(b).unwrap();
650        let collection = rel_target.collection();
651        assert_eq!(collection, &a);
652    }
653
654    #[test]
655    fn one_to_one_relationships() {
656        #[derive(Component)]
657        #[relationship(relationship_target = Below)]
658        struct Above(Entity);
659
660        #[derive(Component)]
661        #[relationship_target(relationship = Above)]
662        struct Below(Entity);
663
664        let mut world = World::new();
665        let a = world.spawn_empty().id();
666        let b = world.spawn_empty().id();
667
668        world.entity_mut(a).insert(Above(b));
669        assert_eq!(a, world.get::<Below>(b).unwrap().0);
670
671        // Verify removing target removes relationship
672        world.entity_mut(b).remove::<Below>();
673        assert!(world.get::<Above>(a).is_none());
674
675        // Verify removing relationship removes target
676        world.entity_mut(a).insert(Above(b));
677        world.entity_mut(a).remove::<Above>();
678        assert!(world.get::<Below>(b).is_none());
679
680        // Actually - a is above c now! Verify relationship was updated correctly
681        let c = world.spawn_empty().id();
682        world.entity_mut(a).insert(Above(c));
683        assert!(world.get::<Below>(b).is_none());
684        assert_eq!(a, world.get::<Below>(c).unwrap().0);
685    }
686
687    #[test]
688    fn entity_index_map() {
689        for add_before in [false, true] {
690            #[derive(Component)]
691            #[relationship(relationship_target = RelTarget)]
692            struct Rel(Entity);
693
694            #[derive(Component)]
695            #[relationship_target(relationship = Rel, linked_spawn)]
696            struct RelTarget(Vec<Entity>);
697
698            let mut world = World::new();
699            if add_before {
700                let _ = world.spawn_empty().id();
701            }
702            let a = world.spawn_empty().id();
703            let b = world.spawn_empty().id();
704            let c = world.spawn_empty().id();
705            let d = world.spawn_empty().id();
706
707            world.entity_mut(a).add_related::<Rel>(&[b, c, d]);
708
709            let rel_target = world.get::<RelTarget>(a).unwrap();
710            let collection = rel_target.collection();
711
712            // Insertions should maintain ordering
713            assert!(collection.iter().eq([b, c, d]));
714
715            world.entity_mut(c).despawn();
716
717            let rel_target = world.get::<RelTarget>(a).unwrap();
718            let collection = rel_target.collection();
719
720            // Removals should maintain ordering
721            assert!(collection.iter().eq([b, d]));
722        }
723    }
724
725    #[test]
726    fn one_to_one_relationship_shared_target() {
727        #[derive(Component)]
728        #[relationship(relationship_target = Below)]
729        struct Above(Entity);
730
731        #[derive(Component)]
732        #[relationship_target(relationship = Above)]
733        struct Below(Entity);
734        let mut world = World::new();
735        let a = world.spawn_empty().id();
736        let b = world.spawn_empty().id();
737        let c = world.spawn_empty().id();
738
739        world.entity_mut(a).insert(Above(c));
740        world.entity_mut(b).insert(Above(c));
741
742        // The original relationship (a -> c) should be removed and the new relationship (b -> c) should be established
743        assert!(
744            world.get::<Above>(a).is_none(),
745            "Original relationship should be removed"
746        );
747        assert_eq!(
748            world.get::<Above>(b).unwrap().0,
749            c,
750            "New relationship should be established"
751        );
752        assert_eq!(
753            world.get::<Below>(c).unwrap().0,
754            b,
755            "Target should point to new source"
756        );
757    }
758
759    #[test]
760    fn one_to_one_relationship_reinsert() {
761        #[derive(Component)]
762        #[relationship(relationship_target = Below)]
763        struct Above(Entity);
764
765        #[derive(Component)]
766        #[relationship_target(relationship = Above)]
767        struct Below(Entity);
768
769        let mut world = World::new();
770        let a = world.spawn_empty().id();
771        let b = world.spawn_empty().id();
772
773        world.entity_mut(a).insert(Above(b));
774        world.entity_mut(a).insert(Above(b));
775    }
776}
bevy_ecs/relationship/relationship_source_collection.rs

bevy_ecs/relationship/
relationship_source_collection.rs
