From f516de456be2b65572de1af6956d52157663b9cc Mon Sep 17 00:00:00 2001 From: Peter Hayman Date: Wed, 3 Apr 2024 20:50:32 +1100 Subject: [PATCH] Add `EntityWorldMut::remove_by_id` (#12842) # Objective - Add `remove_by_id` method to `EntityWorldMut` and `EntityCommands` - This is a duplicate of the awesome work by @mateuseap, last updated 04/09/23 - #9663 - I'm opening a second one to ensure the feature makes it into `0.14` - Fixes #9261 ## Solution Almost identical to #9663 with three exceptions - Uses a closure instead of struct for commands, consistent with other similar commands - Does not refactor `EntityCommands::insert`, so no migration guide - `EntityWorldMut::remove_by_id` is now safe containing unsafe blocks, I think thats what @SkiFire13 was indicating should happen [in this comment](https://github.com/bevyengine/bevy/pull/9663#discussion_r1314307525) ## Changelog - Added `EntityWorldMut::remove_by_id` method and its tests. - Added `EntityCommands::remove_by_id` method and its tests. --------- Co-authored-by: James Liu --- crates/bevy_ecs/src/system/commands/mod.rs | 84 ++++++++++++++++++++-- crates/bevy_ecs/src/world/entity_ref.rs | 37 ++++++++++ 2 files changed, 116 insertions(+), 5 deletions(-) diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index e4fa7f6552..3af9a4e9da 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -4,6 +4,7 @@ use super::{Deferred, IntoSystem, RegisterSystem, Resource}; use crate::{ self as bevy_ecs, bundle::Bundle, + component::ComponentId, entity::{Entities, Entity}, system::{RunSystemWithInput, SystemId}, world::{Command, CommandQueue, EntityWorldMut, FromWorld, World}, @@ -893,6 +894,11 @@ impl EntityCommands<'_> { self.add(remove::) } + /// Removes a component from the entity. + pub fn remove_by_id(&mut self, component_id: ComponentId) -> &mut Self { + self.add(remove_by_id(component_id)) + } + /// Despawns the entity. /// /// See [`World::despawn`] for more details. @@ -1102,8 +1108,20 @@ fn try_insert(bundle: impl Bundle) -> impl EntityCommand { /// For a [`Bundle`] type `T`, this will remove any components in the bundle. /// Any components in the bundle that aren't found on the entity will be ignored. fn remove(entity: Entity, world: &mut World) { - if let Some(mut entity_mut) = world.get_entity_mut(entity) { - entity_mut.remove::(); + if let Some(mut entity) = world.get_entity_mut(entity) { + entity.remove::(); + } +} + +/// An [`EntityCommand`] that removes components with a provided [`ComponentId`] from an entity. +/// # Panics +/// +/// Panics if the provided [`ComponentId`] does not exist in the [`World`]. +fn remove_by_id(component_id: ComponentId) -> impl EntityCommand { + move |entity: Entity, world: &mut World| { + if let Some(mut entity) = world.get_entity_mut(entity) { + entity.remove_by_id(component_id); + } } } @@ -1153,9 +1171,12 @@ mod tests { system::{Commands, Resource}, world::{CommandQueue, World}, }; - use std::sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, + use std::{ + any::TypeId, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, }; #[allow(dead_code)] @@ -1282,6 +1303,59 @@ mod tests { assert_eq!(results_after_u64, vec![]); } + #[test] + fn remove_components_by_id() { + let mut world = World::default(); + + let mut command_queue = CommandQueue::default(); + let (dense_dropck, dense_is_dropped) = DropCk::new_pair(); + let (sparse_dropck, sparse_is_dropped) = DropCk::new_pair(); + let sparse_dropck = SparseDropCk(sparse_dropck); + + let entity = Commands::new(&mut command_queue, &world) + .spawn((W(1u32), W(2u64), dense_dropck, sparse_dropck)) + .id(); + command_queue.apply(&mut world); + let results_before = world + .query::<(&W, &W)>() + .iter(&world) + .map(|(a, b)| (a.0, b.0)) + .collect::>(); + assert_eq!(results_before, vec![(1u32, 2u64)]); + + // test component removal + Commands::new(&mut command_queue, &world) + .entity(entity) + .remove_by_id(world.components().get_id(TypeId::of::>()).unwrap()) + .remove_by_id(world.components().get_id(TypeId::of::>()).unwrap()) + .remove_by_id(world.components().get_id(TypeId::of::()).unwrap()) + .remove_by_id( + world + .components() + .get_id(TypeId::of::()) + .unwrap(), + ); + + assert_eq!(dense_is_dropped.load(Ordering::Relaxed), 0); + assert_eq!(sparse_is_dropped.load(Ordering::Relaxed), 0); + command_queue.apply(&mut world); + assert_eq!(dense_is_dropped.load(Ordering::Relaxed), 1); + assert_eq!(sparse_is_dropped.load(Ordering::Relaxed), 1); + + let results_after = world + .query::<(&W, &W)>() + .iter(&world) + .map(|(a, b)| (a.0, b.0)) + .collect::>(); + assert_eq!(results_after, vec![]); + let results_after_u64 = world + .query::<&W>() + .iter(&world) + .map(|v| v.0) + .collect::>(); + assert_eq!(results_after_u64, vec![]); + } + #[test] fn remove_resources() { let mut world = World::default(); diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index d307de9249..3c79e13e0b 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -1151,6 +1151,27 @@ impl<'w> EntityWorldMut<'w> { self } + /// Removes a dynamic [`Component`] from the entity if it exists. + /// + /// You should prefer to use the typed API [`EntityWorldMut::remove`] where possible. + /// + /// # Panics + /// + /// Panics if the provided [`ComponentId`] does not exist in the [`World`]. + pub fn remove_by_id(&mut self, component_id: ComponentId) -> &mut Self { + let components = &mut self.world.components; + + let bundle_id = self + .world + .bundles + .init_component_info(components, component_id); + + // SAFETY: the `BundleInfo` for this `component_id` is initialized above + self.location = unsafe { self.remove_bundle(bundle_id) }; + + self + } + /// Despawns the current entity. /// /// See [`World::despawn`] for more details. @@ -2767,6 +2788,22 @@ mod tests { assert_eq!(dynamic_components, static_components); } + #[test] + fn entity_mut_remove_by_id() { + let mut world = World::new(); + let test_component_id = world.init_component::(); + + let mut entity = world.spawn(TestComponent(42)); + entity.remove_by_id(test_component_id); + + let components: Vec<_> = world.query::<&TestComponent>().iter(&world).collect(); + + assert_eq!(components, vec![] as Vec<&TestComponent>); + + // remove non-existent component does not panic + world.spawn_empty().remove_by_id(test_component_id); + } + #[derive(Component)] struct A;