Finish #17558, re-adding insert_children (#18409)

fixes #17478

# Objective

- Complete #17558.
- the `insert_children` method was previously removed, and as #17478
points out, needs to be added back.

## Solution

- Add a `OrderedRelationshipSourceCollection`, which allows sorting,
ordering, rearranging, etc of a `RelationshipSourceCollection`.
- Implement `insert_related`
- Implement `insert_children`
- Tidy up some docs while I'm here.

## Testing

@bjoernp116 set up a unit test, and I added a doc test to
`OrderedRelationshipSourceCollection`.

---------

Co-authored-by: bjoernp116 <bjoernpollen@gmail.com>
Co-authored-by: Dmytro Banin <banind@cs.washington.edu>
Co-authored-by: Talin <viridia@gmail.com>

crates/bevy_ecs/src/hierarchy.rs

@@ -256,17 +256,26 @@ pub type ChildSpawnerCommands<'w> = RelatedSpawnerCommands<'w, ChildOf>;
 
 impl<'w> EntityWorldMut<'w> {
     /// Spawns children of this entity (with a [`ChildOf`] relationship) by taking a function that operates on a [`ChildSpawner`].
+    /// See also [`with_related`](Self::with_related).
     pub fn with_children(&mut self, func: impl FnOnce(&mut ChildSpawner)) -> &mut Self {
         self.with_related(func);
         self
     }
 
     /// Adds the given children to this entity
+    /// See also [`add_related`](Self::add_related).
     pub fn add_children(&mut self, children: &[Entity]) -> &mut Self {
         self.add_related::<ChildOf>(children)
     }
 
+    /// Insert children at specific index.
+    /// See also [`insert_related`](Self::insert_related).
+    pub fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self {
+        self.insert_related::<ChildOf>(index, children)
+    }
+
     /// Adds the given child to this entity
+    /// See also [`add_related`](Self::add_related).
     pub fn add_child(&mut self, child: Entity) -> &mut Self {
         self.add_related::<ChildOf>(&[child])
     }
@@ -596,6 +605,35 @@ mod tests {
         );
     }
 
+    #[test]
+    fn insert_children() {
+        let mut world = World::new();
+        let child1 = world.spawn_empty().id();
+        let child2 = world.spawn_empty().id();
+        let child3 = world.spawn_empty().id();
+        let child4 = world.spawn_empty().id();
+
+        let mut entity_world_mut = world.spawn_empty();
+
+        let first_children = entity_world_mut.add_children(&[child1, child2]);
+
+        let root = first_children.insert_children(1, &[child3, child4]).id();
+
+        let hierarchy = get_hierarchy(&world, root);
+        assert_eq!(
+            hierarchy,
+            Node::new_with(
+                root,
+                vec![
+                    Node::new(child1),
+                    Node::new(child3),
+                    Node::new(child4),
+                    Node::new(child2)
+                ]
+            )
+        );
+    }
+
     #[test]
     fn self_parenting_invalid() {
         let mut world = World::new();

crates/bevy_ecs/src/relationship/mod.rs

@@ -209,12 +209,14 @@ pub trait RelationshipTarget: Component<Mutability = Mutable> + Sized {
     ///
     /// # Warning
     /// This should generally not be called by user code, as modifying the internal collection could invalidate the relationship.
+    /// The collection should not contain duplicates.
     fn collection_mut_risky(&mut self) -> &mut Self::Collection;
 
     /// Creates a new [`RelationshipTarget`] from the given [`RelationshipTarget::Collection`].
     ///
     /// # Warning
     /// This should generally not be called by user code, as constructing the internal collection could invalidate the relationship.
+    /// The collection should not contain duplicates.
     fn from_collection_risky(collection: Self::Collection) -> Self;
 
     /// The `on_replace` component hook that maintains the [`Relationship`] / [`RelationshipTarget`] connection.

crates/bevy_ecs/src/relationship/related_methods.rs

@@ -10,6 +10,8 @@ use crate::{
 use bevy_platform_support::prelude::{Box, Vec};
 use core::{marker::PhantomData, mem};
 
+use super::OrderedRelationshipSourceCollection;
+
 impl<'w> EntityWorldMut<'w> {
     /// Spawns entities related to this entity (with the `R` relationship) by taking a function that operates on a [`RelatedSpawner`].
     pub fn with_related<R: Relationship>(
@@ -36,6 +38,64 @@ impl<'w> EntityWorldMut<'w> {
         self
     }
 
+    /// Relates the given entities to this entity with the relation `R`, starting at this particular index.
+    ///
+    /// If the `related` has duplicates, a related entity will take the index of its last occurrence in `related`.
+    /// If the indices go out of bounds, they will be clamped into bounds.
+    /// This will not re-order existing related entities unless they are in `related`.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    ///
+    /// let mut world = World::new();
+    /// let e0 = world.spawn_empty().id();
+    /// let e1 = world.spawn_empty().id();
+    /// let e2 = world.spawn_empty().id();
+    /// let e3 = world.spawn_empty().id();
+    /// let e4 = world.spawn_empty().id();
+    ///
+    /// let mut main_entity = world.spawn_empty();
+    /// main_entity.add_related::<ChildOf>(&[e0, e1, e2, e2]);
+    /// main_entity.insert_related::<ChildOf>(1, &[e0, e3, e4, e4]);
+    /// let main_id = main_entity.id();
+    ///
+    /// let relationship_source = main_entity.get::<Children>().unwrap().collection();
+    /// assert_eq!(relationship_source, &[e1, e0, e3, e2, e4]);
+    /// ```
+    pub fn insert_related<R: Relationship>(&mut self, index: usize, related: &[Entity]) -> &mut Self
+    where
+        <R::RelationshipTarget as RelationshipTarget>::Collection:
+            OrderedRelationshipSourceCollection,
+    {
+        let id = self.id();
+        self.world_scope(|world| {
+            for (offset, related) in related.iter().enumerate() {
+                let index = index + offset;
+                if world
+                    .get::<R>(*related)
+                    .is_some_and(|relationship| relationship.get() == id)
+                {
+                    world
+                        .get_mut::<R::RelationshipTarget>(id)
+                        .expect("hooks should have added relationship target")
+                        .collection_mut_risky()
+                        .place(*related, index);
+                } else {
+                    world.entity_mut(*related).insert(R::from(id));
+                    world
+                        .get_mut::<R::RelationshipTarget>(id)
+                        .expect("hooks should have added relationship target")
+                        .collection_mut_risky()
+                        .place_most_recent(index);
+                }
+            }
+        });
+
+        self
+    }
+
     /// Replaces all the related entities with a new set of entities.
     pub fn replace_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
         type Collection<R> =

crates/bevy_ecs/src/relationship/relationship_source_collection.rs

@@ -74,6 +74,58 @@ pub trait RelationshipSourceCollection {
     }
 }
 
+/// This trait signals that a [`RelationshipSourceCollection`] is ordered.
+pub trait OrderedRelationshipSourceCollection: RelationshipSourceCollection {
+    /// Inserts the entity at a specific index.
+    /// If the index is too large, the entity will be added to the end of the collection.
+    fn insert(&mut self, index: usize, entity: Entity);
+    /// Removes the entity at the specified idnex if it exists.
+    fn remove_at(&mut self, index: usize) -> Option<Entity>;
+    /// Inserts the entity at a specific index.
+    /// This will never reorder other entities.
+    /// If the index is too large, the entity will be added to the end of the collection.
+    fn insert_stable(&mut self, index: usize, entity: Entity);
+    /// Removes the entity at the specified idnex if it exists.
+    /// This will never reorder other entities.
+    fn remove_at_stable(&mut self, index: usize) -> Option<Entity>;
+    /// Sorts the source collection.
+    fn sort(&mut self);
+    /// Inserts the entity at the proper place to maintain sorting.
+    fn insert_sorted(&mut self, entity: Entity);
+
+    /// This places the most recently added entity at the particular index.
+    fn place_most_recent(&mut self, index: usize);
+
+    /// This places the given entity at the particular index.
+    /// This will do nothing if the entity is not in the collection.
+    /// If the index is out of bounds, this will put the entity at the end.
+    fn place(&mut self, entity: Entity, index: usize);
+
+    /// Adds the entity at index 0.
+    fn push_front(&mut self, entity: Entity) {
+        self.insert(0, entity);
+    }
+
+    /// Adds the entity to the back of the collection.
+    fn push_back(&mut self, entity: Entity) {
+        self.insert(usize::MAX, entity);
+    }
+
+    /// Removes the first entity.
+    fn pop_front(&mut self) -> Option<Entity> {
+        self.remove_at(0)
+    }
+
+    /// Removes the last entity.
+    fn pop_back(&mut self) -> Option<Entity> {
+        if self.is_empty() {
+            None
+        } else {
+            self.remove_at(self.len() - 1)
+        }
+    }
+}
+
 impl RelationshipSourceCollection for Vec<Entity> {
     type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
 
@@ -98,7 +150,6 @@ impl RelationshipSourceCollection for Vec<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;
         }
 
@@ -126,6 +177,57 @@ impl RelationshipSourceCollection for Vec<Entity> {
     }
 }
 
+impl OrderedRelationshipSourceCollection for Vec<Entity> {
+    fn insert(&mut self, index: usize, entity: Entity) {
+        self.push(entity);
+        let len = self.len();
+        if index < len {
+            self.swap(index, len - 1);
+        }
+    }
+
+    fn remove_at(&mut self, index: usize) -> Option<Entity> {
+        (index < self.len()).then(|| self.swap_remove(index))
+    }
+
+    fn insert_stable(&mut self, index: usize, entity: Entity) {
+        if index < self.len() {
+            Vec::insert(self, index, entity);
+        } else {
+            self.push(entity);
+        }
+    }
+
+    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
+        (index < self.len()).then(|| self.remove(index))
+    }
+
+    fn sort(&mut self) {
+        self.sort_unstable();
+    }
+
+    fn insert_sorted(&mut self, entity: Entity) {
+        let index = self.partition_point(|e| e <= &entity);
+        self.insert_stable(index, entity);
+    }
+
+    fn place_most_recent(&mut self, index: usize) {
+        if let Some(entity) = self.pop() {
+            let index = index.min(self.len() - 1);
+            self.insert(index, entity);
+        }
+    }
+
+    fn place(&mut self, entity: Entity, index: usize) {
+        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
+            // The len is at least 1, so the subtraction is safe.
+            let index = index.min(self.len() - 1);
+            Vec::remove(self, current);
+            self.insert(index, entity);
+        };
+    }
+}
+
 impl RelationshipSourceCollection for EntityHashSet {
     type SourceIter<'a> = core::iter::Copied<crate::entity::hash_set::Iter<'a>>;
 
@@ -196,7 +298,6 @@ impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
     fn remove(&mut self, entity: Entity) -> bool {
         if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
             SmallVec::remove(self, index);
-
             return true;
         }
 
@@ -277,6 +378,57 @@ impl RelationshipSourceCollection for Entity {
     }
 }
 
+impl<const N: usize> OrderedRelationshipSourceCollection for SmallVec<[Entity; N]> {
+    fn insert(&mut self, index: usize, entity: Entity) {
+        self.push(entity);
+        let len = self.len();
+        if index < len {
+            self.swap(index, len - 1);
+        }
+    }
+
+    fn remove_at(&mut self, index: usize) -> Option<Entity> {
+        (index < self.len()).then(|| self.swap_remove(index))
+    }
+
+    fn insert_stable(&mut self, index: usize, entity: Entity) {
+        if index < self.len() {
+            SmallVec::<[Entity; N]>::insert(self, index, entity);
+        } else {
+            self.push(entity);
+        }
+    }
+
+    fn remove_at_stable(&mut self, index: usize) -> Option<Entity> {
+        (index < self.len()).then(|| self.remove(index))
+    }
+
+    fn sort(&mut self) {
+        self.sort_unstable();
+    }
+
+    fn insert_sorted(&mut self, entity: Entity) {
+        let index = self.partition_point(|e| e <= &entity);
+        self.insert_stable(index, entity);
+    }
+
+    fn place_most_recent(&mut self, index: usize) {
+        if let Some(entity) = self.pop() {
+            let index = index.min(self.len() - 1);
+            self.insert(index, entity);
+        }
+    }
+
+    fn place(&mut self, entity: Entity, index: usize) {
+        if let Some(current) = <[Entity]>::iter(self).position(|e| *e == entity) {
+            // The len is at least 1, so the subtraction is safe.
+            let index = index.min(self.len() - 1);
+            SmallVec::<[Entity; N]>::remove(self, current);
+            self.insert(index, entity);
+        };
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;