From 0dede8e4b48abab804000e9ac43d6a5d9742f19c Mon Sep 17 00:00:00 2001
From: Brezak <bezak.adamko@gmail.com>
Date: Mon, 10 Mar 2025 20:46:20 +0100
Subject: [PATCH] Make the `RelationshipSourceCollection` api report the
 success of `add` and `remove`

---
 .../relationship_source_collection.rs         | 68 +++++++++++++++----
 1 file changed, 56 insertions(+), 12 deletions(-)

diff --git a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs
index e6068bca49a2b..c3748e46a7eab 100644
--- a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs
+++ b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs
@@ -20,10 +20,17 @@ pub trait RelationshipSourceCollection {
     fn with_capacity(capacity: usize) -> Self;
 
     /// Adds the given `entity` to the collection.
-    fn add(&mut self, entity: Entity);
+    ///
+    /// Returns whether the entity was added to the collection.
+    /// Mainly useful when dealing with collections that don't allow
+    /// multiple instances of the same entity ([`EntityHashSet`]).
+    fn add(&mut self, entity: Entity) -> bool;
 
     /// Removes the given `entity` from the collection.
-    fn remove(&mut self, entity: Entity);
+    ///
+    /// Returns whether the collection actually contained
+    /// the entity.
+    fn remove(&mut self, entity: Entity) -> bool;
 
     /// Iterates all entities in the collection.
     fn iter(&self) -> Self::SourceIter<'_>;
@@ -31,6 +38,9 @@ pub trait RelationshipSourceCollection {
     /// Returns the current length of the collection.
     fn len(&self) -> usize;
 
+    /// Clears the collection.
+    fn clear(&mut self);
+
     /// Returns true if the collection contains no entities.
     #[inline]
     fn is_empty(&self) -> bool {
@@ -45,14 +55,20 @@ impl RelationshipSourceCollection for Vec<Entity> {
         Vec::with_capacity(capacity)
     }
 
-    fn add(&mut self, entity: Entity) {
+    fn add(&mut self, entity: Entity) -> bool {
         Vec::push(self, entity);
+
+        true
     }
 
-    fn remove(&mut self, entity: Entity) {
+    fn remove(&mut self, entity: Entity) -> bool {
         if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
             Vec::remove(self, index);
+
+            return true;
         }
+
+        false
     }
 
     fn iter(&self) -> Self::SourceIter<'_> {
@@ -62,6 +78,10 @@ impl RelationshipSourceCollection for Vec<Entity> {
     fn len(&self) -> usize {
         Vec::len(self)
     }
+
+    fn clear(&mut self) {
+        self.clear();
+    }
 }
 
 impl RelationshipSourceCollection for EntityHashSet {
@@ -71,14 +91,14 @@ impl RelationshipSourceCollection for EntityHashSet {
         EntityHashSet::with_capacity(capacity)
     }
 
-    fn add(&mut self, entity: Entity) {
-        self.insert(entity);
+    fn add(&mut self, entity: Entity) -> bool {
+        self.insert(entity)
     }
 
-    fn remove(&mut self, entity: Entity) {
+    fn remove(&mut self, entity: Entity) -> bool {
         // We need to call the remove method on the underlying hash set,
         // which takes its argument by reference
-        self.0.remove(&entity);
+        self.0.remove(&entity)
     }
 
     fn iter(&self) -> Self::SourceIter<'_> {
@@ -88,6 +108,10 @@ impl RelationshipSourceCollection for EntityHashSet {
     fn len(&self) -> usize {
         self.len()
     }
+
+    fn clear(&mut self) {
+        self.0.clear();
+    }
 }
 
 impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
@@ -97,14 +121,20 @@ impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
         SmallVec::with_capacity(capacity)
     }
 
-    fn add(&mut self, entity: Entity) {
+    fn add(&mut self, entity: Entity) -> bool {
         SmallVec::push(self, entity);
+
+        true
     }
 
-    fn remove(&mut self, entity: Entity) {
+    fn remove(&mut self, entity: Entity) -> bool {
         if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
             SmallVec::remove(self, index);
+
+            return true;
         }
+
+        false
     }
 
     fn iter(&self) -> Self::SourceIter<'_> {
@@ -114,6 +144,10 @@ impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
     fn len(&self) -> usize {
         SmallVec::len(self)
     }
+
+    fn clear(&mut self) {
+        self.clear();
+    }
 }
 
 impl RelationshipSourceCollection for Entity {
@@ -123,14 +157,20 @@ impl RelationshipSourceCollection for Entity {
         Entity::PLACEHOLDER
     }
 
-    fn add(&mut self, entity: Entity) {
+    fn add(&mut self, entity: Entity) -> bool {
         *self = entity;
+
+        true
     }
 
-    fn remove(&mut self, entity: Entity) {
+    fn remove(&mut self, entity: Entity) -> bool {
         if *self == entity {
             *self = Entity::PLACEHOLDER;
+
+            return true;
         }
+
+        false
     }
 
     fn iter(&self) -> Self::SourceIter<'_> {
@@ -143,6 +183,10 @@ impl RelationshipSourceCollection for Entity {
         }
         1
     }
+
+    fn clear(&mut self) {
+        *self = Entity::PLACEHOLDER;
+    }
 }
 
 #[cfg(test)]