bevy_ecs/relationship/
relationship_source_collection.rs

1use crate::entity::{hash_set::EntityHashSet, Entity};
2use alloc::vec::Vec;
3use smallvec::SmallVec;
4
5/// The internal [`Entity`] collection used by a [`RelationshipTarget`](crate::relationship::RelationshipTarget) component.
6/// This is not intended to be modified directly by users, as it could invalidate the correctness of relationships.
7pub trait RelationshipSourceCollection {
8    /// The type of iterator returned by the `iter` method.
9    ///
10    /// This is an associated type (rather than using a method that returns an opaque return-position impl trait)
11    /// to ensure that all methods and traits (like [`DoubleEndedIterator`]) of the underlying collection's iterator
12    /// are available to the user when implemented without unduly restricting the possible collections.
13    ///
14    /// The [`SourceIter`](super::SourceIter) type alias can be helpful to reduce confusion when working with this associated type.
15    type SourceIter<'a>: Iterator<Item = Entity>
16    where
17        Self: 'a;
18
19    /// Creates a new empty instance.
20    fn new() -> Self;
21
22    /// Returns an instance with the given pre-allocated entity `capacity`.
23    ///
24    /// Some collections will ignore the provided `capacity` and return a default instance.
25    fn with_capacity(capacity: usize) -> Self;
26
27    /// Reserves capacity for at least `additional` more entities to be inserted.
28    ///
29    /// Not all collections support this operation, in which case it is a no-op.
30    fn reserve(&mut self, additional: usize);
31
32    /// Adds the given `entity` to the collection.
33    ///
34    /// Returns whether the entity was added to the collection.
35    /// Mainly useful when dealing with collections that don't allow
36    /// multiple instances of the same entity ([`EntityHashSet`]).
37    fn add(&mut self, entity: Entity) -> bool;
38
39    /// Removes the given `entity` from the collection.
40    ///
41    /// Returns whether the collection actually contained
42    /// the entity.
43    fn remove(&mut self, entity: Entity) -> bool;
44
45    /// Iterates all entities in the collection.
46    fn iter(&self) -> Self::SourceIter<'_>;
47
48    /// Returns the current length of the collection.
49    fn len(&self) -> usize;
50
51    /// Clears the collection.
52    fn clear(&mut self);
53
54    /// Attempts to save memory by shrinking the capacity to fit the current length.
55    ///
56    /// This operation is a no-op for collections that do not support it.
57    fn shrink_to_fit(&mut self);
58
59    /// Returns true if the collection contains no entities.
60    #[inline]
61    fn is_empty(&self) -> bool {
62        self.len() == 0
63    }
64
65    /// Add multiple entities to collection at once.
66    ///
67    /// May be faster than repeatedly calling [`Self::add`].
68    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
69        // The method name shouldn't conflict with `Extend::extend` as it's in the rust prelude and
70        // would always conflict with it.
71        for entity in entities {
72            self.add(entity);
73        }
74    }
75}
76
77/// This trait signals that a [`RelationshipSourceCollection`] is ordered.
78pub trait OrderedRelationshipSourceCollection: RelationshipSourceCollection {
79    /// Inserts the entity at a specific index.
80    /// If the index is too large, the entity will be added to the end of the collection.
81    fn insert(&mut self, index: usize, entity: Entity);
82    /// Removes the entity at the specified idnex if it exists.
83    fn remove_at(&mut self, index: usize) -> Option<Entity>;
84    /// Inserts the entity at a specific index.
85    /// This will never reorder other entities.
86    /// If the index is too large, the entity will be added to the end of the collection.
87    fn insert_stable(&mut self, index: usize, entity: Entity);
88    /// Removes the entity at the specified idnex if it exists.
89    /// This will never reorder other entities.
90    fn remove_at_stable(&mut self, index: usize) -> Option<Entity>;
91    /// Sorts the source collection.
92    fn sort(&mut self);
93    /// Inserts the entity at the proper place to maintain sorting.
94    fn insert_sorted(&mut self, entity: Entity);
95
96    /// This places the most recently added entity at the particular index.
97    fn place_most_recent(&mut self, index: usize);
98
99    /// This places the given entity at the particular index.
100    /// This will do nothing if the entity is not in the collection.
101    /// If the index is out of bounds, this will put the entity at the end.
102    fn place(&mut self, entity: Entity, index: usize);
103
104    /// Adds the entity at index 0.
105    fn push_front(&mut self, entity: Entity) {
106        self.insert(0, entity);
107    }
108
109    /// Adds the entity to the back of the collection.
110    fn push_back(&mut self, entity: Entity) {
111        self.insert(usize::MAX, entity);
112    }
113
114    /// Removes the first entity.
115    fn pop_front(&mut self) -> Option<Entity> {
116        self.remove_at(0)
117    }
118
119    /// Removes the last entity.
120    fn pop_back(&mut self) -> Option<Entity> {
121        if self.is_empty() {
122            None
123        } else {
124            self.remove_at(self.len() - 1)
125        }
126    }
127}
128
129impl RelationshipSourceCollection for Vec<Entity> {
130    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
131
132    fn new() -> Self {
133        Vec::new()
134    }
135
136    fn reserve(&mut self, additional: usize) {
137        Vec::reserve(self, additional);
138    }
139
140    fn with_capacity(capacity: usize) -> Self {
141        Vec::with_capacity(capacity)
142    }
143
144    fn add(&mut self, entity: Entity) -> bool {
145        Vec::push(self, entity);
146
147        true
148    }
149
150    fn remove(&mut self, entity: Entity) -> bool {
151        if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
152            Vec::remove(self, index);
153            return true;
154        }
155
156        false
157    }
158
159    fn iter(&self) -> Self::SourceIter<'_> {
160        <[Entity]>::iter(self).copied()
161    }
162
163    fn len(&self) -> usize {
164        Vec::len(self)
165    }
166
167    fn clear(&mut self) {
168        self.clear();
169    }
170
171    fn shrink_to_fit(&mut self) {
172        Vec::shrink_to_fit(self);
173    }
174
175    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
176        self.extend(entities);
177    }
178}
179
180impl OrderedRelationshipSourceCollection for Vec<Entity> {
181    fn insert(&mut self, index: usize, entity: Entity) {
182        self.push(entity);
183        let len = self.len();
184        if index < len {
185            self.swap(index, len - 1);
186        }
187    }
188
189    fn remove_at(&mut self, index: usize) -> Option<Entity> {
190        (index < self.len()).then(|| self.swap_remove(index))
191    }
192
193    fn insert_stable(&mut self, index: usize, entity: Entity) {
194        if index < self.len() {
195            Vec::insert(self, index, entity);
196        } else {
197            self.push(entity);
198        }
199    }
200
201    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
202        (index < self.len()).then(|| self.remove(index))
203    }
204
205    fn sort(&mut self) {
206        self.sort_unstable();
207    }
208
209    fn insert_sorted(&mut self, entity: Entity) {
210        let index = self.partition_point(|e| e <= &entity);
211        self.insert_stable(index, entity);
212    }
213
214    fn place_most_recent(&mut self, index: usize) {
215        if let Some(entity) = self.pop() {
216            let index = index.min(self.len().saturating_sub(1));
217            self.insert(index, entity);
218        }
219    }
220
221    fn place(&mut self, entity: Entity, index: usize) {
222        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
223            // The len is at least 1, so the subtraction is safe.
224            let index = index.min(self.len().saturating_sub(1));
225            Vec::remove(self, current);
226            self.insert(index, entity);
227        };
228    }
229}
230
231impl RelationshipSourceCollection for EntityHashSet {
232    type SourceIter<'a> = core::iter::Copied<crate::entity::hash_set::Iter<'a>>;
233
234    fn new() -> Self {
235        EntityHashSet::new()
236    }
237
238    fn reserve(&mut self, additional: usize) {
239        self.0.reserve(additional);
240    }
241
242    fn with_capacity(capacity: usize) -> Self {
243        EntityHashSet::with_capacity(capacity)
244    }
245
246    fn add(&mut self, entity: Entity) -> bool {
247        self.insert(entity)
248    }
249
250    fn remove(&mut self, entity: Entity) -> bool {
251        // We need to call the remove method on the underlying hash set,
252        // which takes its argument by reference
253        self.0.remove(&entity)
254    }
255
256    fn iter(&self) -> Self::SourceIter<'_> {
257        self.iter().copied()
258    }
259
260    fn len(&self) -> usize {
261        self.len()
262    }
263
264    fn clear(&mut self) {
265        self.0.clear();
266    }
267
268    fn shrink_to_fit(&mut self) {
269        self.0.shrink_to_fit();
270    }
271
272    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
273        self.extend(entities);
274    }
275}
276
277impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
278    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
279
280    fn new() -> Self {
281        SmallVec::new()
282    }
283
284    fn reserve(&mut self, additional: usize) {
285        SmallVec::reserve(self, additional);
286    }
287
288    fn with_capacity(capacity: usize) -> Self {
289        SmallVec::with_capacity(capacity)
290    }
291
292    fn add(&mut self, entity: Entity) -> bool {
293        SmallVec::push(self, entity);
294
295        true
296    }
297
298    fn remove(&mut self, entity: Entity) -> bool {
299        if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
300            SmallVec::remove(self, index);
301            return true;
302        }
303
304        false
305    }
306
307    fn iter(&self) -> Self::SourceIter<'_> {
308        <[Entity]>::iter(self).copied()
309    }
310
311    fn len(&self) -> usize {
312        SmallVec::len(self)
313    }
314
315    fn clear(&mut self) {
316        self.clear();
317    }
318
319    fn shrink_to_fit(&mut self) {
320        SmallVec::shrink_to_fit(self);
321    }
322
323    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
324        self.extend(entities);
325    }
326}
327
328impl RelationshipSourceCollection for Entity {
329    type SourceIter<'a> = core::option::IntoIter<Entity>;
330
331    fn new() -> Self {
332        Entity::PLACEHOLDER
333    }
334
335    fn reserve(&mut self, _: usize) {}
336
337    fn with_capacity(_capacity: usize) -> Self {
338        Self::new()
339    }
340
341    fn add(&mut self, entity: Entity) -> bool {
342        assert_eq!(
343            *self,
344            Entity::PLACEHOLDER,
345            "Entity {entity} attempted to target an entity with a one-to-one relationship, but it is already targeted by {}. You must remove the original relationship first.",
346            *self
347        );
348        *self = entity;
349
350        true
351    }
352
353    fn remove(&mut self, entity: Entity) -> bool {
354        if *self == entity {
355            *self = Entity::PLACEHOLDER;
356
357            return true;
358        }
359
360        false
361    }
362
363    fn iter(&self) -> Self::SourceIter<'_> {
364        if *self == Entity::PLACEHOLDER {
365            None.into_iter()
366        } else {
367            Some(*self).into_iter()
368        }
369    }
370
371    fn len(&self) -> usize {
372        if *self == Entity::PLACEHOLDER {
373            return 0;
374        }
375        1
376    }
377
378    fn clear(&mut self) {
379        *self = Entity::PLACEHOLDER;
380    }
381
382    fn shrink_to_fit(&mut self) {}
383
384    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
385        for entity in entities {
386            assert_eq!(
387                *self,
388                Entity::PLACEHOLDER,
389                "Entity {entity} attempted to target an entity with a one-to-one relationship, but it is already targeted by {}. You must remove the original relationship first.",
390                *self
391            );
392            *self = entity;
393        }
394    }
395}
396
397impl<const N: usize> OrderedRelationshipSourceCollection for SmallVec<[Entity; N]> {
398    fn insert(&mut self, index: usize, entity: Entity) {
399        self.push(entity);
400        let len = self.len();
401        if index < len {
402            self.swap(index, len - 1);
403        }
404    }
405
406    fn remove_at(&mut self, index: usize) -> Option<Entity> {
407        (index < self.len()).then(|| self.swap_remove(index))
408    }
409
410    fn insert_stable(&mut self, index: usize, entity: Entity) {
411        if index < self.len() {
412            SmallVec::<[Entity; N]>::insert(self, index, entity);
413        } else {
414            self.push(entity);
415        }
416    }
417
418    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
419        (index < self.len()).then(|| self.remove(index))
420    }
421
422    fn sort(&mut self) {
423        self.sort_unstable();
424    }
425
426    fn insert_sorted(&mut self, entity: Entity) {
427        let index = self.partition_point(|e| e <= &entity);
428        self.insert_stable(index, entity);
429    }
430
431    fn place_most_recent(&mut self, index: usize) {
432        if let Some(entity) = self.pop() {
433            let index = index.min(self.len() - 1);
434            self.insert(index, entity);
435        }
436    }
437
438    fn place(&mut self, entity: Entity, index: usize) {
439        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
440            // The len is at least 1, so the subtraction is safe.
441            let index = index.min(self.len() - 1);
442            SmallVec::<[Entity; N]>::remove(self, current);
443            self.insert(index, entity);
444        };
445    }
446}
447
448#[cfg(test)]
449mod tests {
450    use super::*;
451    use crate::prelude::{Component, World};
452    use crate::relationship::RelationshipTarget;
453
454    #[test]
455    fn vec_relationship_source_collection() {
456        #[derive(Component)]
457        #[relationship(relationship_target = RelTarget)]
458        struct Rel(Entity);
459
460        #[derive(Component)]
461        #[relationship_target(relationship = Rel, linked_spawn)]
462        struct RelTarget(Vec<Entity>);
463
464        let mut world = World::new();
465        let a = world.spawn_empty().id();
466        let b = world.spawn_empty().id();
467
468        world.entity_mut(a).insert(Rel(b));
469
470        let rel_target = world.get::<RelTarget>(b).unwrap();
471        let collection = rel_target.collection();
472        assert_eq!(collection, &alloc::vec!(a));
473    }
474
475    #[test]
476    fn smallvec_relationship_source_collection() {
477        #[derive(Component)]
478        #[relationship(relationship_target = RelTarget)]
479        struct Rel(Entity);
480
481        #[derive(Component)]
482        #[relationship_target(relationship = Rel, linked_spawn)]
483        struct RelTarget(SmallVec<[Entity; 4]>);
484
485        let mut world = World::new();
486        let a = world.spawn_empty().id();
487        let b = world.spawn_empty().id();
488
489        world.entity_mut(a).insert(Rel(b));
490
491        let rel_target = world.get::<RelTarget>(b).unwrap();
492        let collection = rel_target.collection();
493        assert_eq!(collection, &SmallVec::from_buf([a]));
494    }
495
496    #[test]
497    fn entity_relationship_source_collection() {
498        #[derive(Component)]
499        #[relationship(relationship_target = RelTarget)]
500        struct Rel(Entity);
501
502        #[derive(Component)]
503        #[relationship_target(relationship = Rel)]
504        struct RelTarget(Entity);
505
506        let mut world = World::new();
507        let a = world.spawn_empty().id();
508        let b = world.spawn_empty().id();
509
510        world.entity_mut(a).insert(Rel(b));
511
512        let rel_target = world.get::<RelTarget>(b).unwrap();
513        let collection = rel_target.collection();
514        assert_eq!(collection, &a);
515    }
516
517    #[test]
518    fn one_to_one_relationships() {
519        #[derive(Component)]
520        #[relationship(relationship_target = Below)]
521        struct Above(Entity);
522
523        #[derive(Component)]
524        #[relationship_target(relationship = Above)]
525        struct Below(Entity);
526
527        let mut world = World::new();
528        let a = world.spawn_empty().id();
529        let b = world.spawn_empty().id();
530
531        world.entity_mut(a).insert(Above(b));
532        assert_eq!(a, world.get::<Below>(b).unwrap().0);
533
534        // Verify removing target removes relationship
535        world.entity_mut(b).remove::<Below>();
536        assert!(world.get::<Above>(a).is_none());
537
538        // Verify removing relationship removes target
539        world.entity_mut(a).insert(Above(b));
540        world.entity_mut(a).remove::<Above>();
541        assert!(world.get::<Below>(b).is_none());
542
543        // Actually - a is above c now! Verify relationship was updated correctly
544        let c = world.spawn_empty().id();
545        world.entity_mut(a).insert(Above(c));
546        assert!(world.get::<Below>(b).is_none());
547        assert_eq!(a, world.get::<Below>(c).unwrap().0);
548    }
549
550    #[test]
551    #[should_panic]
552    fn one_to_one_relationship_shared_target() {
553        #[derive(Component)]
554        #[relationship(relationship_target = Below)]
555        struct Above(Entity);
556
557        #[derive(Component)]
558        #[relationship_target(relationship = Above)]
559        struct Below(Entity);
560
561        let mut world = World::new();
562        let a = world.spawn_empty().id();
563        let b = world.spawn_empty().id();
564        let c = world.spawn_empty().id();
565
566        world.entity_mut(a).insert(Above(c));
567        world.entity_mut(b).insert(Above(c));
568    }
569
570    #[test]
571    fn one_to_one_relationship_reinsert() {
572        #[derive(Component)]
573        #[relationship(relationship_target = Below)]
574        struct Above(Entity);
575
576        #[derive(Component)]
577        #[relationship_target(relationship = Above)]
578        struct Below(Entity);
579
580        let mut world = World::new();
581        let a = world.spawn_empty().id();
582        let b = world.spawn_empty().id();
583
584        world.entity_mut(a).insert(Above(b));
585        world.entity_mut(a).insert(Above(b));
586    }
587}
bevy_ecs/relationship/relationship_source_collection.rs

bevy_ecs/relationship/
relationship_source_collection.rs
