Add more methods to RelationshipSourceCollection

crates/bevy_ecs/src/relationship/relationship_source_collection.rs

@@ -16,9 +16,19 @@ pub trait RelationshipSourceCollection {
     where
         Self: 'a;
 
+    /// Creates a new empty instance.
+    fn new() -> Self;
+
     /// Returns an instance with the given pre-allocated entity `capacity`.
+    ///
+    /// Some collections will ignore the provided `capacity` and return a default instance.
     fn with_capacity(capacity: usize) -> Self;
 
+    /// Reserves capacity for at least `additional` more entities to be inserted.
+    ///
+    /// Not all collections support this operation, in which case it is a no-op.
+    fn reserve(&mut self, additional: usize);
+
     /// Adds the given `entity` to the collection.
     ///
     /// Returns whether the entity was added to the collection.
@@ -41,6 +51,11 @@ pub trait RelationshipSourceCollection {
     /// Clears the collection.
     fn clear(&mut self);
 
+    /// Attempts to save memory by shrinking the capacity to fit the current length.
+    ///
+    /// This operation is a no-op for collections that do not support it.
+    fn shrink_to_fit(&mut self);
+
     /// Returns true if the collection contains no entities.
     #[inline]
     fn is_empty(&self) -> bool {
@@ -62,6 +77,14 @@ pub trait RelationshipSourceCollection {
 impl RelationshipSourceCollection for Vec<Entity> {
     type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
 
+    fn new() -> Self {
+        Vec::new()
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        Vec::reserve(self, additional);
+    }
+
     fn with_capacity(capacity: usize) -> Self {
         Vec::with_capacity(capacity)
     }
@@ -94,6 +117,10 @@ impl RelationshipSourceCollection for Vec<Entity> {
         self.clear();
     }
 
+    fn shrink_to_fit(&mut self) {
+        Vec::shrink_to_fit(self);
+    }
+
     fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
         self.extend(entities);
     }
@@ -102,6 +129,14 @@ impl RelationshipSourceCollection for Vec<Entity> {
 impl RelationshipSourceCollection for EntityHashSet {
     type SourceIter<'a> = core::iter::Copied<crate::entity::hash_set::Iter<'a>>;
 
+    fn new() -> Self {
+        EntityHashSet::new()
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        self.0.reserve(additional);
+    }
+
     fn with_capacity(capacity: usize) -> Self {
         EntityHashSet::with_capacity(capacity)
     }
@@ -128,6 +163,10 @@ impl RelationshipSourceCollection for EntityHashSet {
         self.0.clear();
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.0.shrink_to_fit();
+    }
+
     fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
         self.extend(entities);
     }
@@ -136,6 +175,14 @@ impl RelationshipSourceCollection for EntityHashSet {
 impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
     type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
 
+    fn new() -> Self {
+        SmallVec::new()
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        SmallVec::reserve(self, additional);
+    }
+
     fn with_capacity(capacity: usize) -> Self {
         SmallVec::with_capacity(capacity)
     }
@@ -168,6 +215,10 @@ impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
         self.clear();
     }
 
+    fn shrink_to_fit(&mut self) {
+        SmallVec::shrink_to_fit(self);
+    }
+
     fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
         self.extend(entities);
     }
@@ -176,10 +227,16 @@ impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
 impl RelationshipSourceCollection for Entity {
     type SourceIter<'a> = core::iter::Once<Entity>;
 
-    fn with_capacity(_capacity: usize) -> Self {
+    fn new() -> Self {
         Entity::PLACEHOLDER
     }
 
+    fn reserve(&mut self, _: usize) {}
+
+    fn with_capacity(_capacity: usize) -> Self {
+        Self::new()
+    }
+
     fn add(&mut self, entity: Entity) -> bool {
         *self = entity;
 
@@ -211,6 +268,8 @@ impl RelationshipSourceCollection for Entity {
         *self = Entity::PLACEHOLDER;
     }
 
+    fn shrink_to_fit(&mut self) {}
+
     fn extend_from_iter(&mut self, entities: impl IntoIterator<Item = Entity>) {
         if let Some(entity) = entities.into_iter().last() {
             *self = entity;