From 0f446d49a1cc7bd5ded209bc44807bfb4b604558 Mon Sep 17 00:00:00 2001
From: Brezak <bezak.adam@proton.me>
Date: Thu, 20 Mar 2025 02:35:51 +0100
Subject: [PATCH] Add more methods to `RelationshipSourceCollection` (#18421)

# Objective

While working on #18058 I realized I could use
`RelationshipTargetCollection::new`, so I added it.

## Solution

- Add `RelationshipTargetCollection::new`
- Add `RelationshipTargetCollection::reserve`. Could generally be useful
when doing micro-optimizations.
- Add `RelationshipTargetCollection::shrink_to_fit`. Rust collections
generally don't shrink when removing elements. Might be a good idea to
call this once in a while.

## Testing

`cargo clippy`

---

## Showcase

`RelationshipSourceCollection` now implements `new`, `reserve` and
`shrink_to_fit` to give greater control over how much memory it
consumes.

## Migration Guide

Any type implementing `RelationshipSourceCollection` now needs to also
implement `new`, `reserve` and `shrink_to_fit`. `reserve` and
`shrink_to_fit` can be made no-ops if they conceptually mean nothing to
a collection.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Chris Russell <8494645+chescock@users.noreply.github.com>
---
 .../relationship_source_collection.rs         | 61 ++++++++++++++++++-
 1 file changed, 60 insertions(+), 1 deletion(-)

diff --git a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs
index 5db6851176..2f052c7d88 100644
--- a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs
+++ b/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;