From 9da65b10b4dbe7de4a6ea64f17e0f281dac958b3 Mon Sep 17 00:00:00 2001 From: TheBigCheese <32036861+13ros27@users.noreply.github.com> Date: Tue, 5 Dec 2023 15:37:33 +0000 Subject: [PATCH] Add `EntityCommands.retain` and `EntityWorldMut.retain` (#10873) # Objective Adds `EntityCommands.retain` and `EntityWorldMut.retain` to remove all components except the given bundle from the entity. Fixes #10865. ## Solution I added a private unsafe function in `EntityWorldMut` called `remove_bundle_info` which performs the shared behaviour of `remove` and `retain`, namely taking a `BundleInfo` of components to remove, and removing them from the given entity. Then `retain` simply gets all the components on the entity and filters them by whether they are in the bundle it was passed, before passing this `BundleInfo` into `remove_bundle_info`. `EntityCommands.retain` just creates a new type `Retain` which runs `EntityWorldMut.retain` when run. --- ## Changelog Added `EntityCommands.retain` and `EntityWorldMut.retain`, which remove all components except the given bundle from the entity, they can also be used to remove all components by passing `()` as the bundle. --- crates/bevy_ecs/src/system/commands/mod.rs | 79 ++++++++++ crates/bevy_ecs/src/world/entity_ref.rs | 171 +++++++++++++++++---- 2 files changed, 218 insertions(+), 32 deletions(-) diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index 4b9f8121bd7c6..5f081804da557 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -908,6 +908,54 @@ impl<'w, 's, 'a> EntityCommands<'w, 's, 'a> { self } + /// Removes all components except the given [`Bundle`] from the entity. + /// + /// This can also be used to remove all the components from the entity by passing it an empty Bundle. + /// + /// See [`EntityWorldMut::retain`](EntityWorldMut::retain) for more + /// details. + /// + /// # Example + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Resource)] + /// # struct PlayerEntity { entity: Entity } + /// #[derive(Component)] + /// struct Health(u32); + /// #[derive(Component)] + /// struct Strength(u32); + /// #[derive(Component)] + /// struct Defense(u32); + /// + /// #[derive(Bundle)] + /// struct CombatBundle { + /// health: Health, + /// strength: Strength, + /// } + /// + /// fn remove_combat_stats_system(mut commands: Commands, player: Res) { + /// commands + /// .entity(player.entity) + /// // You can retain a pre-defined Bundle of components, + /// // with this removing only the Defense component + /// .retain::() + /// // You can also retain only a single component + /// .retain::() + /// // And you can remove all the components by passing in an empty Bundle + /// .retain::<()>(); + /// } + /// # bevy_ecs::system::assert_is_system(remove_combat_stats_system); + /// ``` + pub fn retain(&mut self) -> &mut Self + where + T: Bundle, + { + self.commands.add(Retain::::new(self.entity)); + self + } + /// Logs the components of the entity at the info level. /// /// # Panics @@ -1097,6 +1145,37 @@ impl Remove { } } +/// A [`Command`] that removes components from an entity. +/// For a [`Bundle`] type `T`, this will remove all components except those in the bundle. +/// Any components in the bundle that aren't found on the entity will be ignored. +#[derive(Debug)] +pub struct Retain { + /// The entity from which the components will be removed. + pub entity: Entity, + _marker: PhantomData, +} + +impl Command for Retain +where + T: Bundle, +{ + fn apply(self, world: &mut World) { + if let Some(mut entity_mut) = world.get_entity_mut(self.entity) { + entity_mut.retain::(); + } + } +} + +impl Retain { + /// Creates a [`Command`] which will remove all but the specified components when applied. + pub const fn new(entity: Entity) -> Self { + Self { + entity, + _marker: PhantomData, + } + } +} + /// A [`Command`] that inserts a [`Resource`] into the world using a value /// created with the [`FromWorld`] trait. pub struct InitResource { diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index c94b45dad10ee..5f17409c9036f 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -825,38 +825,39 @@ impl<'w> EntityWorldMut<'w> { entities.set(entity.index(), new_location); } - /// Removes any components in the [`Bundle`] from the entity. - // TODO: BundleRemover? - pub fn remove(&mut self) -> &mut Self { - let archetypes = &mut self.world.archetypes; - let storages = &mut self.world.storages; - let components = &mut self.world.components; - let entities = &mut self.world.entities; - let removed_components = &mut self.world.removed_components; - - let bundle_info = self.world.bundles.init_info::(components, storages); - let old_location = self.location; - - // SAFETY: `archetype_id` exists because it is referenced in the old `EntityLocation` which is valid, - // components exist in `bundle_info` because `Bundles::init_info` initializes a `BundleInfo` containing all components of the bundle type `T` - let new_archetype_id = unsafe { - remove_bundle_from_archetype( - archetypes, - storages, - components, - old_location.archetype_id, - bundle_info, - true, - ) - .expect("intersections should always return a result") - }; + /// Remove the components of `bundle_info` from `entity`, where `self_location` and `old_location` + /// are the location of this entity, and `self_location` is updated to the new location. + /// + /// SAFETY: `old_location` must be valid and the components in `bundle_info` must exist. + #[allow(clippy::too_many_arguments)] + unsafe fn remove_bundle_info( + entity: Entity, + self_location: &mut EntityLocation, + old_location: EntityLocation, + bundle_info: &BundleInfo, + archetypes: &mut Archetypes, + storages: &mut Storages, + components: &Components, + entities: &mut Entities, + removed_components: &mut RemovedComponentEvents, + ) { + // SAFETY: `archetype_id` exists because it is referenced in `old_location` which is valid + // and components in `bundle_info` must exist due to this functions safety invariants. + let new_archetype_id = remove_bundle_from_archetype( + archetypes, + storages, + components, + old_location.archetype_id, + bundle_info, + true, + ) + .expect("intersections should always return a result"); if new_archetype_id == old_location.archetype_id { - return self; + return; } let old_archetype = &mut archetypes[old_location.archetype_id]; - let entity = self.entity; for component_id in bundle_info.components().iter().cloned() { if old_archetype.contains(component_id) { removed_components.send(component_id, entity); @@ -873,17 +874,86 @@ impl<'w> EntityWorldMut<'w> { } } - #[allow(clippy::undocumented_unsafe_blocks)] // TODO: document why this is safe + // SAFETY: `new_archetype_id` is a subset of the components in `old_location.archetype_id` + // because it is created by removing a bundle from these components. + Self::move_entity_from_remove::( + entity, + self_location, + old_location.archetype_id, + old_location, + entities, + archetypes, + storages, + new_archetype_id, + ); + } + + /// Removes any components in the [`Bundle`] from the entity. + // TODO: BundleRemover? + pub fn remove(&mut self) -> &mut Self { + let archetypes = &mut self.world.archetypes; + let storages = &mut self.world.storages; + let components = &mut self.world.components; + let entities = &mut self.world.entities; + let removed_components = &mut self.world.removed_components; + + let bundle_info = self.world.bundles.init_info::(components, storages); + let old_location = self.location; + + // SAFETY: Components exist in `bundle_info` because `Bundles::init_info` + // initializes a `BundleInfo` containing all components of the bundle type `T`. unsafe { - Self::move_entity_from_remove::( - entity, + Self::remove_bundle_info( + self.entity, &mut self.location, - old_location.archetype_id, old_location, + bundle_info, + archetypes, + storages, + components, entities, + removed_components, + ); + } + + self + } + + /// Removes any components except those in the [`Bundle`] from the entity. + pub fn retain(&mut self) -> &mut Self { + let archetypes = &mut self.world.archetypes; + let storages = &mut self.world.storages; + let components = &mut self.world.components; + let entities = &mut self.world.entities; + let removed_components = &mut self.world.removed_components; + + let retained_bundle_info = self.world.bundles.init_info::(components, storages); + let old_location = self.location; + let old_archetype = &mut archetypes[old_location.archetype_id]; + + let to_remove = &old_archetype + .components() + .filter(|c| !retained_bundle_info.components().contains(c)) + .collect::>(); + let remove_bundle_info = self + .world + .bundles + .init_dynamic_info(components, to_remove) + .0; + + // SAFETY: Components exist in `remove_bundle_info` because `Bundles::init_dynamic_info` + // initializes a `BundleInfo` containing all components in the to_remove Bundle. + unsafe { + Self::remove_bundle_info( + self.entity, + &mut self.location, + old_location, + remove_bundle_info, archetypes, storages, - new_archetype_id, + components, + entities, + removed_components, ); } @@ -1775,6 +1845,43 @@ mod tests { assert_eq!(world.entity(e2).get::().unwrap(), &Dense(1)); } + // Test that calling retain with `()` removes all components. + #[test] + fn retain_nothing() { + #[derive(Component)] + struct Marker; + + let mut world = World::new(); + let ent = world.spawn((Marker::<1>, Marker::<2>, Marker::<3>)).id(); + + world.entity_mut(ent).retain::<()>(); + assert_eq!(world.entity(ent).archetype().components().next(), None); + } + + // Test removing some components with `retain`, including components not on the entity. + #[test] + fn retain_some_components() { + #[derive(Component)] + struct Marker; + + let mut world = World::new(); + let ent = world.spawn((Marker::<1>, Marker::<2>, Marker::<3>)).id(); + + world.entity_mut(ent).retain::<(Marker<2>, Marker<4>)>(); + // Check that marker 2 was retained. + assert!(world.entity(ent).get::>().is_some()); + // Check that only marker 2 was retained. + assert_eq!( + world + .entity(ent) + .archetype() + .components() + .collect::>() + .len(), + 1 + ); + } + // regression test for https://github.com/bevyengine/bevy/pull/7805 #[test] fn inserting_sparse_updates_archetype_row() {