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());
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            let index = index.min(self.len());
224            Vec::remove(self, current);
225            self.insert(index, entity);
226        };
227    }
228}
229
230impl RelationshipSourceCollection for EntityHashSet {
231    type SourceIter<'a> = core::iter::Copied<crate::entity::hash_set::Iter<'a>>;
232
233    fn new() -> Self {
234        EntityHashSet::new()
235    }
236
237    fn reserve(&mut self, additional: usize) {
238        self.0.reserve(additional);
239    }
240
241    fn with_capacity(capacity: usize) -> Self {
242        EntityHashSet::with_capacity(capacity)
243    }
244
245    fn add(&mut self, entity: Entity) -> bool {
246        self.insert(entity)
247    }
248
249    fn remove(&mut self, entity: Entity) -> bool {
250        // We need to call the remove method on the underlying hash set,
251        // which takes its argument by reference
252        self.0.remove(&entity)
253    }
254
255    fn iter(&self) -> Self::SourceIter<'_> {
256        self.iter().copied()
257    }
258
259    fn len(&self) -> usize {
260        self.len()
261    }
262
263    fn clear(&mut self) {
264        self.0.clear();
265    }
266
267    fn shrink_to_fit(&mut self) {
268        self.0.shrink_to_fit();
269    }
270
271    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
272        self.extend(entities);
273    }
274}
275
276impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
277    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
278
279    fn new() -> Self {
280        SmallVec::new()
281    }
282
283    fn reserve(&mut self, additional: usize) {
284        SmallVec::reserve(self, additional);
285    }
286
287    fn with_capacity(capacity: usize) -> Self {
288        SmallVec::with_capacity(capacity)
289    }
290
291    fn add(&mut self, entity: Entity) -> bool {
292        SmallVec::push(self, entity);
293
294        true
295    }
296
297    fn remove(&mut self, entity: Entity) -> bool {
298        if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
299            SmallVec::remove(self, index);
300            return true;
301        }
302
303        false
304    }
305
306    fn iter(&self) -> Self::SourceIter<'_> {
307        <[Entity]>::iter(self).copied()
308    }
309
310    fn len(&self) -> usize {
311        SmallVec::len(self)
312    }
313
314    fn clear(&mut self) {
315        self.clear();
316    }
317
318    fn shrink_to_fit(&mut self) {
319        SmallVec::shrink_to_fit(self);
320    }
321
322    fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
323        self.extend(entities);
324    }
325}
326
327impl RelationshipSourceCollection for Entity {
328    type SourceIter<'a> = core::option::IntoIter<Entity>;
329
330    fn new() -> Self {
331        Entity::PLACEHOLDER
332    }
333
334    fn reserve(&mut self, _: usize) {}
335
336    fn with_capacity(_capacity: usize) -> Self {
337        Self::new()
338    }
339
340    fn add(&mut self, entity: Entity) -> bool {
341        assert_eq!(
342            *self,
343            Entity::PLACEHOLDER,
344            "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.",
345            *self
346        );
347        *self = entity;
348
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            assert_eq!(
386                *self,
387                Entity::PLACEHOLDER,
388                "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.",
389                *self
390            );
391            *self = entity;
392        }
393    }
394}
395
396impl<const N: usize> OrderedRelationshipSourceCollection for SmallVec<[Entity; N]> {
397    fn insert(&mut self, index: usize, entity: Entity) {
398        self.push(entity);
399        let len = self.len();
400        if index < len {
401            self.swap(index, len - 1);
402        }
403    }
404
405    fn remove_at(&mut self, index: usize) -> Option<Entity> {
406        (index < self.len()).then(|| self.swap_remove(index))
407    }
408
409    fn insert_stable(&mut self, index: usize, entity: Entity) {
410        if index < self.len() {
411            SmallVec::<[Entity; N]>::insert(self, index, entity);
412        } else {
413            self.push(entity);
414        }
415    }
416
417    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
418        (index < self.len()).then(|| self.remove(index))
419    }
420
421    fn sort(&mut self) {
422        self.sort_unstable();
423    }
424
425    fn insert_sorted(&mut self, entity: Entity) {
426        let index = self.partition_point(|e| e <= &entity);
427        self.insert_stable(index, entity);
428    }
429
430    fn place_most_recent(&mut self, index: usize) {
431        if let Some(entity) = self.pop() {
432            let index = index.min(self.len() - 1);
433            self.insert(index, entity);
434        }
435    }
436
437    fn place(&mut self, entity: Entity, index: usize) {
438        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
439            // The len is at least 1, so the subtraction is safe.
440            let index = index.min(self.len() - 1);
441            SmallVec::<[Entity; N]>::remove(self, current);
442            self.insert(index, entity);
443        };
444    }
445}
446
447#[cfg(test)]
448mod tests {
449    use super::*;
450    use crate::prelude::{Component, World};
451    use crate::relationship::RelationshipTarget;
452
453    #[test]
454    fn vec_relationship_source_collection() {
455        #[derive(Component)]
456        #[relationship(relationship_target = RelTarget)]
457        struct Rel(Entity);
458
459        #[derive(Component)]
460        #[relationship_target(relationship = Rel, linked_spawn)]
461        struct RelTarget(Vec<Entity>);
462
463        let mut world = World::new();
464        let a = world.spawn_empty().id();
465        let b = world.spawn_empty().id();
466
467        world.entity_mut(a).insert(Rel(b));
468
469        let rel_target = world.get::<RelTarget>(b).unwrap();
470        let collection = rel_target.collection();
471        assert_eq!(collection, &alloc::vec!(a));
472    }
473
474    #[test]
475    fn smallvec_relationship_source_collection() {
476        #[derive(Component)]
477        #[relationship(relationship_target = RelTarget)]
478        struct Rel(Entity);
479
480        #[derive(Component)]
481        #[relationship_target(relationship = Rel, linked_spawn)]
482        struct RelTarget(SmallVec<[Entity; 4]>);
483
484        let mut world = World::new();
485        let a = world.spawn_empty().id();
486        let b = world.spawn_empty().id();
487
488        world.entity_mut(a).insert(Rel(b));
489
490        let rel_target = world.get::<RelTarget>(b).unwrap();
491        let collection = rel_target.collection();
492        assert_eq!(collection, &SmallVec::from_buf([a]));
493    }
494
495    #[test]
496    fn entity_relationship_source_collection() {
497        #[derive(Component)]
498        #[relationship(relationship_target = RelTarget)]
499        struct Rel(Entity);
500
501        #[derive(Component)]
502        #[relationship_target(relationship = Rel)]
503        struct RelTarget(Entity);
504
505        let mut world = World::new();
506        let a = world.spawn_empty().id();
507        let b = world.spawn_empty().id();
508
509        world.entity_mut(a).insert(Rel(b));
510
511        let rel_target = world.get::<RelTarget>(b).unwrap();
512        let collection = rel_target.collection();
513        assert_eq!(collection, &a);
514    }
515
516    #[test]
517    fn one_to_one_relationships() {
518        #[derive(Component)]
519        #[relationship(relationship_target = Below)]
520        struct Above(Entity);
521
522        #[derive(Component)]
523        #[relationship_target(relationship = Above)]
524        struct Below(Entity);
525
526        let mut world = World::new();
527        let a = world.spawn_empty().id();
528        let b = world.spawn_empty().id();
529
530        world.entity_mut(a).insert(Above(b));
531        assert_eq!(a, world.get::<Below>(b).unwrap().0);
532
533        // Verify removing target removes relationship
534        world.entity_mut(b).remove::<Below>();
535        assert!(world.get::<Above>(a).is_none());
536
537        // Verify removing relationship removes target
538        world.entity_mut(a).insert(Above(b));
539        world.entity_mut(a).remove::<Above>();
540        assert!(world.get::<Below>(b).is_none());
541
542        // Actually - a is above c now! Verify relationship was updated correctly
543        let c = world.spawn_empty().id();
544        world.entity_mut(a).insert(Above(c));
545        assert!(world.get::<Below>(b).is_none());
546        assert_eq!(a, world.get::<Below>(c).unwrap().0);
547    }
548
549    #[test]
550    #[should_panic]
551    fn one_to_one_relationship_shared_target() {
552        #[derive(Component)]
553        #[relationship(relationship_target = Below)]
554        struct Above(Entity);
555
556        #[derive(Component)]
557        #[relationship_target(relationship = Above)]
558        struct Below(Entity);
559
560        let mut world = World::new();
561        let a = world.spawn_empty().id();
562        let b = world.spawn_empty().id();
563        let c = world.spawn_empty().id();
564
565        world.entity_mut(a).insert(Above(c));
566        world.entity_mut(b).insert(Above(c));
567    }
568
569    #[test]
570    fn one_to_one_relationship_reinsert() {
571        #[derive(Component)]
572        #[relationship(relationship_target = Below)]
573        struct Above(Entity);
574
575        #[derive(Component)]
576        #[relationship_target(relationship = Above)]
577        struct Below(Entity);
578
579        let mut world = World::new();
580        let a = world.spawn_empty().id();
581        let b = world.spawn_empty().id();
582
583        world.entity_mut(a).insert(Above(b));
584        world.entity_mut(a).insert(Above(b));
585    }
586}
bevy_ecs/relationship/relationship_source_collection.rs

bevy_ecs/relationship/
relationship_source_collection.rs
