Add a faster version of replace_related

Author: Brezak

crates/bevy_ecs/src/hierarchy.rs

@@ -12,7 +12,7 @@ use crate::{
     bundle::Bundle,
     component::{Component, HookContext},
     entity::Entity,
-    relationship::{RelatedSpawner, RelatedSpawnerCommands},
+    relationship::{RelatedSpawner, RelatedSpawnerCommands, Relationship},
     system::EntityCommands,
     world::{DeferredWorld, EntityWorldMut, FromWorld, World},
 };
@@ -179,11 +179,33 @@ impl<'w> EntityWorldMut<'w> {
         self.add_related::<ChildOf>(&[child])
     }
 
-    /// Replaces the children on this entity with a new list of children.
+    /// Replaces all the related children with a new set of children.
     pub fn replace_children(&mut self, children: &[Entity]) -> &mut Self {
         self.replace_related::<ChildOf>(children)
     }
 
+    /// Replaces all the related children with a new set of children.
+    ///
+    /// # Warning
+    /// Not maintaining the function invariants may lead to erratic engine behavior including random crashes.
+    /// See [`Self::replace_related_with_difference`] for invariants.
+    ///
+    /// # Panics
+    ///
+    /// In debug mode when function invariants are broken.
+    pub fn replace_children_with_difference(
+        &mut self,
+        entities_to_unrelate: &[Entity],
+        entities_to_relate: &[Entity],
+        newly_related_entities: &[Entity],
+    ) -> &mut Self {
+        self.replace_related_with_difference::<ChildOf>(
+            entities_to_unrelate,
+            entities_to_relate,
+            newly_related_entities,
+        )
+    }
+
     /// Spawns the passed bundle and adds it to this entity as a child.
     ///
     /// For efficient spawning of multiple children, use [`with_children`].
@@ -240,6 +262,28 @@ impl<'a> EntityCommands<'a> {
         self.replace_related::<ChildOf>(children)
     }
 
+    /// Replaces all the related entities with a new set of entities.
+    ///
+    /// # Warning
+    /// Not maintaining the function invariants may lead to erratic engine behavior including random crashes.
+    /// See [`EntityWorldMut::replace_related_with_difference`] for invariants.
+    ///
+    /// # Panics
+    ///
+    /// In debug mode when function invariants are broken.
+    pub fn replace_children_with_difference<R: Relationship>(
+        &mut self,
+        entities_to_unrelate: &[Entity],
+        entities_to_relate: &[Entity],
+        newly_related_entities: &[Entity],
+    ) -> &mut Self {
+        self.replace_related_with_difference::<R>(
+            entities_to_unrelate,
+            entities_to_relate,
+            newly_related_entities,
+        )
+    }
+
     /// Spawns the passed bundle and adds it to this entity as a child.
     ///
     /// For efficient spawning of multiple children, use [`with_children`].
@@ -336,6 +380,8 @@ macro_rules! children {
 
 #[cfg(test)]
 mod tests {
+    use core::ops::{Deref, DerefMut};
+
     use crate::{
         entity::Entity,
         hierarchy::{ChildOf, Children},
@@ -533,4 +579,169 @@ mod tests {
 
         assert!(world.entity(parent).get::<Children>().is_none());
     }
+
+    #[test]
+    fn insert_same_child_twice() {
+        let mut world = World::new();
+
+        let parent = world.spawn_empty().id();
+        let child = world.spawn_empty().id();
+
+        world.entity_mut(parent).add_child(child);
+        world.entity_mut(parent).add_child(child);
+
+        let children = world.get::<Children>(parent).unwrap();
+        assert_eq!(children.0, [child]);
+    }
+
+    #[test]
+    fn replace_with_difference() {
+        let mut world = World::new();
+
+        let parent = world.spawn_empty().id();
+        let child_a = world.spawn_empty().id();
+        let child_b = world.spawn_empty().id();
+        let child_c = world.spawn_empty().id();
+        let child_d = world.spawn_empty().id();
+
+        // Test inserting new relations
+        world.entity_mut(parent).replace_children_with_difference(
+            &[],
+            &[child_a, child_b],
+            &[child_a, child_b],
+        );
+
+        assert_eq!(
+            world.entity(child_a).get::<ChildOf>().unwrap(),
+            &ChildOf { parent }
+        );
+        assert_eq!(
+            world.entity(child_b).get::<ChildOf>().unwrap(),
+            &ChildOf { parent }
+        );
+        assert_eq!(
+            world.entity(parent).get::<Children>().unwrap().0,
+            [child_a, child_b]
+        );
+
+        // Test replacing relations and changing order
+        world.entity_mut(parent).replace_children_with_difference(
+            &[child_b],
+            &[child_d, child_c, child_a],
+            &[child_c, child_d],
+        );
+        assert_eq!(
+            world.entity(child_a).get::<ChildOf>().unwrap(),
+            &ChildOf { parent }
+        );
+        assert_eq!(
+            world.entity(child_c).get::<ChildOf>().unwrap(),
+            &ChildOf { parent }
+        );
+        assert_eq!(
+            world.entity(child_d).get::<ChildOf>().unwrap(),
+            &ChildOf { parent }
+        );
+        assert_eq!(
+            world.entity(parent).get::<Children>().unwrap().0,
+            [child_d, child_c, child_a]
+        );
+        assert!(!world.entity(child_b).contains::<ChildOf>());
+
+        // Test removing relationships
+        world.entity_mut(parent).replace_children_with_difference(
+            &[child_a, child_d, child_c],
+            &[],
+            &[],
+        );
+        assert!(!world.entity(parent).contains::<Children>());
+        assert!(!world.entity(child_a).contains::<ChildOf>());
+        assert!(!world.entity(child_b).contains::<ChildOf>());
+        assert!(!world.entity(child_c).contains::<ChildOf>());
+        assert!(!world.entity(child_d).contains::<ChildOf>());
+    }
+
+    #[test]
+    fn replace_with_difference_on_empty() {
+        let mut world = World::new();
+
+        let parent = world.spawn_empty().id();
+        let child_a = world.spawn_empty().id();
+
+        world
+            .entity_mut(parent)
+            .replace_children_with_difference(&[child_a], &[], &[]);
+
+        assert!(!world.entity(parent).contains::<Children>());
+        assert!(!world.entity(child_a).contains::<ChildOf>());
+    }
+
+    #[test]
+    #[should_panic]
+    #[cfg_attr(
+        not(debug_assertions),
+        ignore = "we don't check invariants if debug assertions are off"
+    )]
+    fn replace_diff_invariant_overlapping_unrelate_with_relate() {
+        let mut world = World::new();
+
+        let parent = world.spawn_empty().id();
+        let child_a = world.spawn_empty().id();
+
+        world
+            .entity_mut(parent)
+            .replace_children_with_difference(&[], &[child_a], &[child_a]);
+
+        // This should panic
+        world
+            .entity_mut(parent)
+            .replace_children_with_difference(&[child_a], &[child_a], &[]);
+    }
+
+    #[test]
+    #[should_panic]
+    #[cfg_attr(
+        not(debug_assertions),
+        ignore = "we don't check invariants if debug assertions are off"
+    )]
+
+    fn replace_diff_invariant_overlapping_unrelate_with_newly() {
+        let mut world = World::new();
+
+        let parent = world.spawn_empty().id();
+        let child_a = world.spawn_empty().id();
+        let child_b = world.spawn_empty().id();
+
+        world
+            .entity_mut(parent)
+            .replace_children_with_difference(&[], &[child_a], &[child_a]);
+
+        // This should panic
+        world.entity_mut(parent).replace_children_with_difference(
+            &[child_b],
+            &[child_a, child_b],
+            &[child_b],
+        );
+    }
+
+    #[test]
+    #[should_panic]
+    #[cfg_attr(
+        not(debug_assertions),
+        ignore = "we don't check invariants if debug assertions are off"
+    )]
+    fn replace_diff_invariant_newly_not_subset() {
+        let mut world = World::new();
+
+        let parent = world.spawn_empty().id();
+        let child_a = world.spawn_empty().id();
+        let child_b = world.spawn_empty().id();
+
+        // This should panic
+        world.entity_mut(parent).replace_children_with_difference(
+            &[],
+            &[child_a, child_b],
+            &[child_a],
+        );
+    }
 }