From 28907ae171971073ed646b3ae0e7f14f72eb8e48 Mon Sep 17 00:00:00 2001 From: Brezak Date: Wed, 19 Mar 2025 21:04:42 +0100 Subject: [PATCH] Add methods to bulk replace relationships on a entity (#18058) # Objective Add a way to efficiently replace a set of specifically related entities with a new set. Closes #18041 ## Solution Add new `replace_(related/children)` to `EntityWorldMut` and friends. ## Testing Added a new test to `hierarchy.rs` that specifically check if `replace_children` actually correctly replaces the children on a entity while keeping the original one. --- ## Showcase `EntityWorldMut` and `EntityCommands` can now be used to efficiently replace the entities a entity is related to. ```rust /// `parent` has 2 children. `entity_a` and `entity_b`. assert_eq!([entity_a, entity_b], world.entity(parent).get::()); /// Replace `parent`s children with `entity_a` and `entity_c` world.entity_mut(parent).replace_related(&[entity_a, entity_c]); /// `parent` now has 2 children. `entity_a` and `entity_c`. /// /// `replace_children` has saved time by not removing and reading /// the relationship between `entity_a` and `parent` assert_eq!([entity_a, entity_c], world.entity(parent).get::()); --------- Co-authored-by: Alice Cecile --- crates/bevy_ecs/src/hierarchy.rs | 354 +++++++++++++++++- .../src/relationship/related_methods.rs | 207 +++++++++- .../relationship_source_collection.rs | 29 ++ 3 files changed, 583 insertions(+), 7 deletions(-) diff --git a/crates/bevy_ecs/src/hierarchy.rs b/crates/bevy_ecs/src/hierarchy.rs index 7f1d0c4180..276746f50b 100644 --- a/crates/bevy_ecs/src/hierarchy.rs +++ b/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}, }; @@ -26,7 +26,7 @@ use log::warn; /// Stores the parent entity of this child entity with this component. /// -/// This is a [`Relationship`](crate::relationship::Relationship) component, and creates the canonical +/// This is a [`Relationship`] component, and creates the canonical /// "parent / child" hierarchy. This is the "source of truth" component, and it pairs with /// the [`Children`] [`RelationshipTarget`](crate::relationship::RelationshipTarget). /// @@ -117,7 +117,7 @@ impl FromWorld for ChildOf { /// Tracks which entities are children of this parent entity. /// /// A [`RelationshipTarget`] collection component that is populated -/// with entities that "target" this entity with the [`ChildOf`] [`Relationship`](crate::relationship::Relationship) component. +/// with entities that "target" this entity with the [`ChildOf`] [`Relationship`] component. /// /// Together, these components form the "canonical parent-child hierarchy". See the [`ChildOf`] component for the full /// description of this relationship and instructions on how to use it. @@ -181,6 +181,34 @@ impl<'w> EntityWorldMut<'w> { self.add_related::(&[child]) } + /// Replaces all the related children with a new set of children. + pub fn replace_children(&mut self, children: &[Entity]) -> &mut Self { + self.replace_related::(children) + } + + /// Replaces all the related children with a new set of children. + /// + /// # Warning + /// + /// Failing to maintain the functions invariants may lead to erratic engine behavior including random crashes. + /// Refer to [`Self::replace_related_with_difference`] for a list of these invariants. + /// + /// # Panics + /// + /// Panics when debug assertions are enabled if an invariant is is broken and the command is executed. + 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::( + 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`]. @@ -232,6 +260,34 @@ impl<'a> EntityCommands<'a> { self.add_related::(&[child]) } + /// Replaces the children on this entity with a new list of children. + pub fn replace_children(&mut self, children: &[Entity]) -> &mut Self { + self.replace_related::(children) + } + + /// Replaces all the related entities with a new set of entities. + /// + /// # Warning + /// + /// Failing to maintain the functions invariants may lead to erratic engine behavior including random crashes. + /// Refer to [`EntityWorldMut::replace_related_with_difference`] for a list of these invariants. + /// + /// # Panics + /// + /// Panics when debug assertions are enabled if an invariant is is broken and the command is executed. + 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::( + 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`]. @@ -493,6 +549,298 @@ mod tests { assert_eq!(world.entity(id).get::().unwrap().len(), 2,); } + #[test] + fn replace_children() { + let mut world = World::new(); + let parent = world.spawn(Children::spawn((Spawn(()), Spawn(())))).id(); + let &[child_a, child_b] = &world.entity(parent).get::().unwrap().0[..] else { + panic!("Tried to spawn 2 children on an entity and didn't get 2 children"); + }; + + let child_c = world.spawn_empty().id(); + + world + .entity_mut(parent) + .replace_children(&[child_a, child_c]); + + let children = world.entity(parent).get::().unwrap(); + + assert!(children.contains(&child_a)); + assert!(children.contains(&child_c)); + assert!(!children.contains(&child_b)); + + assert_eq!( + world.entity(child_a).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(child_c).get::().unwrap(), + &ChildOf { parent } + ); + assert!(world.entity(child_b).get::().is_none()); + } + + #[test] + fn replace_children_with_nothing() { + 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).add_children(&[child_a, child_b]); + + assert_eq!(world.entity(parent).get::().unwrap().len(), 2); + + world.entity_mut(parent).replace_children(&[]); + + assert!(world.entity(child_a).get::().is_none()); + assert!(world.entity(child_b).get::().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::(parent).unwrap(); + assert_eq!(children.0, [child]); + assert_eq!( + world.entity(child).get::().unwrap(), + &ChildOf { parent } + ); + } + + #[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::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(child_b).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(parent).get::().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::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(child_c).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(child_d).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(parent).get::().unwrap().0, + [child_d, child_c, child_a] + ); + assert!(!world.entity(child_b).contains::()); + + // Test removing relationships + world.entity_mut(parent).replace_children_with_difference( + &[child_a, child_d, child_c], + &[], + &[], + ); + assert!(!world.entity(parent).contains::()); + assert!(!world.entity(child_a).contains::()); + assert!(!world.entity(child_b).contains::()); + assert!(!world.entity(child_c).contains::()); + assert!(!world.entity(child_d).contains::()); + } + + #[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::()); + assert!(!world.entity(child_a).contains::()); + } + + #[test] + fn replace_with_difference_totally_new_children() { + 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::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(child_b).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(parent).get::().unwrap().0, + [child_a, child_b] + ); + + // Test replacing relations and changing order + world.entity_mut(parent).replace_children_with_difference( + &[child_b, child_a], + &[child_d, child_c], + &[child_c, child_d], + ); + assert_eq!( + world.entity(child_c).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(child_d).get::().unwrap(), + &ChildOf { parent } + ); + assert_eq!( + world.entity(parent).get::().unwrap().0, + [child_d, child_c] + ); + assert!(!world.entity(child_a).contains::()); + assert!(!world.entity(child_b).contains::()); + } + + #[test] + fn replace_children_order() { + 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(); + + let initial_order = [child_a, child_b, child_c, child_d]; + world.entity_mut(parent).add_children(&initial_order); + + assert_eq!( + world.entity_mut(parent).get::().unwrap().0, + initial_order + ); + + let new_order = [child_d, child_b, child_a, child_c]; + world.entity_mut(parent).replace_children(&new_order); + + assert_eq!(world.entity(parent).get::().unwrap().0, new_order); + } + + #[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], + ); + } + #[test] fn child_replace_hook_skip() { let mut world = World::new(); diff --git a/crates/bevy_ecs/src/relationship/related_methods.rs b/crates/bevy_ecs/src/relationship/related_methods.rs index b13baf5f2a..9d9d8711cb 100644 --- a/crates/bevy_ecs/src/relationship/related_methods.rs +++ b/crates/bevy_ecs/src/relationship/related_methods.rs @@ -1,12 +1,14 @@ use crate::{ bundle::Bundle, - entity::Entity, - relationship::{Relationship, RelationshipTarget}, + entity::{hash_set::EntityHashSet, Entity}, + relationship::{ + Relationship, RelationshipHookMode, RelationshipSourceCollection, RelationshipTarget, + }, system::{Commands, EntityCommands}, world::{EntityWorldMut, World}, }; -use alloc::vec::Vec; -use core::marker::PhantomData; +use bevy_platform_support::prelude::{Box, Vec}; +use core::{marker::PhantomData, mem}; impl<'w> EntityWorldMut<'w> { /// Spawns entities related to this entity (with the `R` relationship) by taking a function that operates on a [`RelatedSpawner`]. @@ -34,6 +36,159 @@ impl<'w> EntityWorldMut<'w> { self } + /// Replaces all the related entities with a new set of entities. + pub fn replace_related(&mut self, related: &[Entity]) -> &mut Self { + type Collection = + <::RelationshipTarget as RelationshipTarget>::Collection; + + if related.is_empty() { + self.remove::(); + + return self; + } + + let Some(mut existing_relations) = self.get_mut::() else { + return self.add_related::(related); + }; + + // We take the collection here so we can modify it without taking the component itself (this would create archetype move). + // SAFETY: We eventually return the correctly initialized collection into the target. + let mut existing_relations = mem::replace( + existing_relations.collection_mut_risky(), + Collection::::with_capacity(0), + ); + + let mut potential_relations = EntityHashSet::from_iter(related.iter().copied()); + + let id = self.id(); + self.world_scope(|world| { + for related in existing_relations.iter() { + if !potential_relations.remove(related) { + world.entity_mut(related).remove::(); + } + } + + for related in potential_relations { + // SAFETY: We'll manually be adjusting the contents of the parent to fit the final state. + world + .entity_mut(related) + .insert_with_relationship_hook_mode(R::from(id), RelationshipHookMode::Skip); + } + }); + + // SAFETY: The entities we're inserting will be the entities that were either already there or entities that we've just inserted. + existing_relations.clear(); + existing_relations.extend_from_iter(related.iter().copied()); + self.insert(R::RelationshipTarget::from_collection_risky( + existing_relations, + )); + + self + } + + /// Replaces all the related entities with a new set of entities. + /// + /// This is a more efficient of [`Self::replace_related`] which doesn't allocate. + /// The passed in arguments must adhere to these invariants: + /// - `entities_to_unrelate`: A slice of entities to remove from the relationship source. + /// Entities need not be related to this entity, but must not appear in `entities_to_relate` + /// - `entities_to_relate`: A slice of entities to relate to this entity. + /// This must contain all entities that will remain related (i.e. not those in `entities_to_unrelate`) plus the newly related entities. + /// - `newly_related_entities`: A subset of `entities_to_relate` containing only entities not already related to this entity. + /// - Slices **must not** contain any duplicates + /// + /// # Warning + /// + /// Violating these invariants may lead to panics, crashes or unpredictable engine behavior. + /// + /// # Panics + /// + /// Panics when debug assertions are enabled and any invariants are broken. + /// + // TODO: Consider making these iterators so users aren't required to allocate a separate buffers for the different slices. + pub fn replace_related_with_difference( + &mut self, + entities_to_unrelate: &[Entity], + entities_to_relate: &[Entity], + newly_related_entities: &[Entity], + ) -> &mut Self { + #[cfg(debug_assertions)] + { + let entities_to_relate = EntityHashSet::from_iter(entities_to_relate.iter().copied()); + let entities_to_unrelate = + EntityHashSet::from_iter(entities_to_unrelate.iter().copied()); + let mut newly_related_entities = + EntityHashSet::from_iter(newly_related_entities.iter().copied()); + assert!( + entities_to_relate.is_disjoint(&entities_to_unrelate), + "`entities_to_relate` ({entities_to_relate:?}) shared entities with `entities_to_unrelate` ({entities_to_unrelate:?})" + ); + assert!( + newly_related_entities.is_disjoint(&entities_to_unrelate), + "`newly_related_entities` ({newly_related_entities:?}) shared entities with `entities_to_unrelate ({entities_to_unrelate:?})`" + ); + assert!( + newly_related_entities.is_subset(&entities_to_relate), + "`newly_related_entities` ({newly_related_entities:?}) wasn't a subset of `entities_to_relate` ({entities_to_relate:?})" + ); + + if let Some(target) = self.get::() { + let existing_relationships: EntityHashSet = target.collection().iter().collect(); + + assert!( + existing_relationships.is_disjoint(&newly_related_entities), + "`newly_related_entities` contains an entity that wouldn't be newly related" + ); + + newly_related_entities.extend(existing_relationships); + newly_related_entities -= &entities_to_unrelate; + } + + assert_eq!(newly_related_entities, entities_to_relate, "`entities_to_relate` ({entities_to_relate:?}) didn't contain all entities that would end up related"); + }; + + if !self.contains::() { + self.add_related::(entities_to_relate); + + return self; + }; + + let this = self.id(); + self.world_scope(|world| { + for unrelate in entities_to_unrelate { + world.entity_mut(*unrelate).remove::(); + } + + for new_relation in newly_related_entities { + // We're changing the target collection manually so don't run the insert hook + world + .entity_mut(*new_relation) + .insert_with_relationship_hook_mode(R::from(this), RelationshipHookMode::Skip); + } + }); + + if !entities_to_relate.is_empty() { + if let Some(mut target) = self.get_mut::() { + // SAFETY: The invariants expected by this function mean we'll only be inserting entities that are already related. + let collection = target.collection_mut_risky(); + collection.clear(); + + collection.extend_from_iter(entities_to_relate.iter().copied()); + } else { + let mut empty = + ::Collection::with_capacity( + entities_to_relate.len(), + ); + empty.extend_from_iter(entities_to_relate.iter().copied()); + + // SAFETY: We've just initialized this collection and we know there's no `RelationshipTarget` on `self` + self.insert(R::RelationshipTarget::from_collection_risky(empty)); + } + } + + self + } + /// Relates the given entity to this with the relation `R`. /// /// See [`add_related`](Self::add_related) if you want to relate more than one entity. @@ -138,6 +293,50 @@ impl<'a> EntityCommands<'a> { self.add_related::(&[entity]) } + /// Replaces all the related entities with the given set of new related entities. + pub fn replace_related(&mut self, related: &[Entity]) -> &mut Self { + let id = self.id(); + let related: Box<[Entity]> = related.into(); + + self.commands().queue(move |world: &mut World| { + world.entity_mut(id).replace_related::(&related); + }); + + self + } + + /// Replaces all the related entities with a new set of entities. + /// + /// # Warning + /// + /// Failing to maintain the functions invariants may lead to erratic engine behavior including random crashes. + /// Refer to [`EntityWorldMut::replace_related_with_difference`] for a list of these invariants. + /// + /// # Panics + /// + /// Panics when debug assertions are enable, an invariant is are broken and the command is executed. + pub fn replace_related_with_difference( + &mut self, + entities_to_unrelate: &[Entity], + entities_to_relate: &[Entity], + newly_related_entities: &[Entity], + ) -> &mut Self { + let id = self.id(); + let entities_to_unrelate: Box<[Entity]> = entities_to_unrelate.into(); + let entities_to_relate: Box<[Entity]> = entities_to_relate.into(); + let newly_related_entities: Box<[Entity]> = newly_related_entities.into(); + + self.commands().queue(move |world: &mut World| { + world.entity_mut(id).replace_children_with_difference( + &entities_to_unrelate, + &entities_to_relate, + &newly_related_entities, + ); + }); + + self + } + /// Despawns entities that relate to this one via the given [`RelationshipTarget`]. /// This entity will not be despawned. pub fn despawn_related(&mut self) -> &mut Self { diff --git a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs index c3748e46a7..5db6851176 100644 --- a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs +++ b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs @@ -46,6 +46,17 @@ pub trait RelationshipSourceCollection { fn is_empty(&self) -> bool { self.len() == 0 } + + /// Add multiple entities to collection at once. + /// + /// May be faster than repeatedly calling [`Self::add`]. + fn extend_from_iter(&mut self, entities: impl IntoIterator) { + // The method name shouldn't conflict with `Extend::extend` as it's in the rust prelude and + // would always conflict with it. + for entity in entities { + self.add(entity); + } + } } impl RelationshipSourceCollection for Vec { @@ -82,6 +93,10 @@ impl RelationshipSourceCollection for Vec { fn clear(&mut self) { self.clear(); } + + fn extend_from_iter(&mut self, entities: impl IntoIterator) { + self.extend(entities); + } } impl RelationshipSourceCollection for EntityHashSet { @@ -112,6 +127,10 @@ impl RelationshipSourceCollection for EntityHashSet { fn clear(&mut self) { self.0.clear(); } + + fn extend_from_iter(&mut self, entities: impl IntoIterator) { + self.extend(entities); + } } impl RelationshipSourceCollection for SmallVec<[Entity; N]> { @@ -148,6 +167,10 @@ impl RelationshipSourceCollection for SmallVec<[Entity; N]> { fn clear(&mut self) { self.clear(); } + + fn extend_from_iter(&mut self, entities: impl IntoIterator) { + self.extend(entities); + } } impl RelationshipSourceCollection for Entity { @@ -187,6 +210,12 @@ impl RelationshipSourceCollection for Entity { fn clear(&mut self) { *self = Entity::PLACEHOLDER; } + + fn extend_from_iter(&mut self, entities: impl IntoIterator) { + if let Some(entity) = entities.into_iter().last() { + *self = entity; + } + } } #[cfg(test)]