diff --git a/crates/bevy_ecs/src/hierarchy.rs b/crates/bevy_ecs/src/hierarchy.rs index 276746f50b093..bc40e3ada2580 100644 --- a/crates/bevy_ecs/src/hierarchy.rs +++ b/crates/bevy_ecs/src/hierarchy.rs @@ -166,17 +166,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::(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::(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::(&[child]) } @@ -506,6 +515,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(); diff --git a/crates/bevy_ecs/src/relationship/mod.rs b/crates/bevy_ecs/src/relationship/mod.rs index 9230d77b3f6ef..13e7d9d26b249 100644 --- a/crates/bevy_ecs/src/relationship/mod.rs +++ b/crates/bevy_ecs/src/relationship/mod.rs @@ -212,12 +212,14 @@ pub trait RelationshipTarget: Component + 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. diff --git a/crates/bevy_ecs/src/relationship/related_methods.rs b/crates/bevy_ecs/src/relationship/related_methods.rs index 9d9d8711cbddf..1717210b2c3b2 100644 --- a/crates/bevy_ecs/src/relationship/related_methods.rs +++ b/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( @@ -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::(&[e0, e1, e2, e2]); + /// main_entity.insert_related::(1, &[e0, e3, e4, e4]); + /// let main_id = main_entity.id(); + /// + /// let relationship_source = main_entity.get::().unwrap().collection(); + /// assert_eq!(relationship_source, &[e1, e0, e3, e2, e4]); + /// ``` + pub fn insert_related(&mut self, index: usize, related: &[Entity]) -> &mut Self + where + ::Collection: + OrderedRelationshipSourceCollection, + { + let id = self.id(); + self.world_scope(|world| { + for (offset, related) in related.iter().enumerate() { + let index = index + offset; + if world + .get::(*related) + .is_some_and(|relationship| relationship.get() == id) + { + world + .get_mut::(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::(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(&mut self, related: &[Entity]) -> &mut Self { type Collection = diff --git a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs index 5db6851176fcd..605e6c1133857 100644 --- a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs +++ b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs @@ -59,6 +59,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; + /// 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; + /// 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 { + self.remove_at(0) + } + + /// Removes the last entity. + fn pop_back(&mut self) -> Option { + if self.is_empty() { + None + } else { + self.remove_at(self.len() - 1) + } + } +} + impl RelationshipSourceCollection for Vec { type SourceIter<'a> = core::iter::Copied>; @@ -75,7 +127,6 @@ impl RelationshipSourceCollection for Vec { fn remove(&mut self, entity: Entity) -> bool { if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) { Vec::remove(self, index); - return true; } @@ -99,6 +150,57 @@ impl RelationshipSourceCollection for Vec { } } +impl OrderedRelationshipSourceCollection for Vec { + 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 { + (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 { + (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>; @@ -149,7 +251,6 @@ impl 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; } @@ -218,6 +319,57 @@ impl RelationshipSourceCollection for Entity { } } +impl 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 { + (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 { + (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::*;