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.

Author: 13ros27

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<PlayerEntity>) {
+    ///     commands
+    ///         .entity(player.entity)
+    ///         // You can retain a pre-defined Bundle of components,
+    ///         // with this removing only the Defense component
+    ///         .retain::<CombatBundle>()
+    ///         // You can also retain only a single component
+    ///         .retain::<Health>()
+    ///         // 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<T>(&mut self) -> &mut Self
+    where
+        T: Bundle,
+    {
+        self.commands.add(Retain::<T>::new(self.entity));
+        self
+    }
+
     /// Logs the components of the entity at the info level.
     ///
     /// # Panics
@@ -1097,6 +1145,37 @@ impl<T> Remove<T> {
     }
 }
 
+/// 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<T> {
+    /// The entity from which the components will be removed.
+    pub entity: Entity,
+    _marker: PhantomData<T>,
+}
+
+impl<T> Command for Retain<T>
+where
+    T: Bundle,
+{
+    fn apply(self, world: &mut World) {
+        if let Some(mut entity_mut) = world.get_entity_mut(self.entity) {
+            entity_mut.retain::<T>();
+        }
+    }
+}
+
+impl<T> Retain<T> {
+    /// 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<R: Resource + FromWorld> {

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<T: Bundle>(&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::<T>(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::<true>(
+            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<T: Bundle>(&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::<T>(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::<true>(
-                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<T: Bundle>(&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::<T>(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::<Vec<_>>();
+        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::<Dense>().unwrap(), &Dense(1));
     }
 
+    // Test that calling retain with `()` removes all components.
+    #[test]
+    fn retain_nothing() {
+        #[derive(Component)]
+        struct Marker<const N: usize>;
+
+        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<const N: usize>;
+
+        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::<Marker<2>>().is_some());
+        // Check that only marker 2 was retained.
+        assert_eq!(
+            world
+                .entity(ent)
+                .archetype()
+                .components()
+                .collect::<Vec<_>>()
+                .len(),
+            1
+        );
+    }
+
     // regression test for https://github.com/bevyengine/bevy/pull/7805
     #[test]
     fn inserting_sparse_updates_archetype_row() {