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 <contact@jamessliu.com>

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::<T>)
     }
 
+    /// 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<T: Bundle>(entity: Entity, world: &mut World) {
-    if let Some(mut entity_mut) = world.get_entity_mut(entity) {
+    if let Some(mut entity) = world.get_entity_mut(entity) {
-        entity_mut.remove::<T>();
+        entity.remove::<T>();
+    }
+}
+
+/// 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::{
+    use std::{
-        atomic::{AtomicUsize, Ordering},
+        any::TypeId,
-        Arc,
+        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<u32>, &W<u64>)>()
+            .iter(&world)
+            .map(|(a, b)| (a.0, b.0))
+            .collect::<Vec<_>>();
+        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::<W<u32>>()).unwrap())
+            .remove_by_id(world.components().get_id(TypeId::of::<W<u64>>()).unwrap())
+            .remove_by_id(world.components().get_id(TypeId::of::<DropCk>()).unwrap())
+            .remove_by_id(
+                world
+                    .components()
+                    .get_id(TypeId::of::<SparseDropCk>())
+                    .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<u32>, &W<u64>)>()
+            .iter(&world)
+            .map(|(a, b)| (a.0, b.0))
+            .collect::<Vec<_>>();
+        assert_eq!(results_after, vec![]);
+        let results_after_u64 = world
+            .query::<&W<u64>>()
+            .iter(&world)
+            .map(|v| v.0)
+            .collect::<Vec<_>>();
+        assert_eq!(results_after_u64, vec![]);
+    }
+
     #[test]
     fn remove_resources() {
         let mut world = World::default();

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::<TestComponent>();
+
+        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;