Add method to remove component and all required components for removed component

crates/bevy_ecs/src/bundle.rs

@@ -1302,6 +1302,8 @@ pub struct Bundles {
     bundle_infos: Vec<BundleInfo>,
     /// Cache static [`BundleId`]
     bundle_ids: TypeIdMap<BundleId>,
+    /// Cache bundles, which contains both explicit and required components of [`Bundle`]
+    contributed_bundle_ids: TypeIdMap<BundleId>,
     /// Cache dynamic [`BundleId`] with multiple components
     dynamic_bundle_ids: HashMap<Box<[ComponentId]>, BundleId>,
     dynamic_bundle_storages: HashMap<BundleId, Vec<StorageType>>,
@@ -1335,6 +1337,7 @@ impl Bundles {
         storages: &mut Storages,
     ) -> BundleId {
         let bundle_infos = &mut self.bundle_infos;
+        let contributed_bundle_ids = &mut self.contributed_bundle_ids;
         let id = *self.bundle_ids.entry(TypeId::of::<T>()).or_insert_with(|| {
             let mut component_ids= Vec::new();
             T::component_ids(components, storages, &mut |id| component_ids.push(id));
@@ -1346,11 +1349,39 @@ impl Bundles {
                 // - it was created in the same order as the components in T
                 unsafe { BundleInfo::new(core::any::type_name::<T>(), components, component_ids, id) };
             bundle_infos.push(bundle_info);
+
+            // Be sure to update component bundle with requires
+            contributed_bundle_ids.remove(&TypeId::of::<T>());
+
             id
         });
         id
     }
 
+    /// Registers a new [`BundleInfo`], which contains both explicit and required components for a statically known type.
+    ///
+    /// Also registers all the components in the bundle.
+    pub(crate) fn register_contributed_bundle_info<T: Bundle>(
+        &mut self,
+        components: &mut Components,
+        storages: &mut Storages,
+    ) -> BundleId {
+        if let Some(id) = self.contributed_bundle_ids.get(&TypeId::of::<T>()).cloned() {
+            id
+        } else {
+            let explicit_bundle_id = self.register_info::<T>(components, storages);
+            // SAFETY: `explicit_bundle_id` is valid and defined above
+            let contributed_components = unsafe {
+                self.get_unchecked(explicit_bundle_id)
+                    .contributed_components()
+                    .to_vec()
+            };
+            let id = self.init_dynamic_info(components, &contributed_components);
+            self.contributed_bundle_ids.insert(TypeId::of::<T>(), id);
+            id
+        }
+    }
+
     /// # Safety
     /// A [`BundleInfo`] with the given [`BundleId`] must have been initialized for this instance of `Bundles`.
     pub(crate) unsafe fn get_unchecked(&self, id: BundleId) -> &BundleInfo {

crates/bevy_ecs/src/lib.rs

@@ -2029,6 +2029,138 @@ mod tests {
         assert!(e.contains::<Y>());
     }
 
+    #[test]
+    fn remove_component_and_his_runtime_required_components() {
+        #[derive(Component)]
+        struct X;
+
+        #[derive(Component, Default)]
+        struct Y;
+
+        #[derive(Component, Default)]
+        struct Z;
+
+        #[derive(Component)]
+        struct V;
+
+        let mut world = World::new();
+        world.register_required_components::<X, Y>();
+        world.register_required_components::<Y, Z>();
+
+        let e = world.spawn((X, V)).id();
+        assert!(world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+
+        //check that `remove` works as expected
+        world.entity_mut(e).remove::<X>();
+        assert!(!world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+
+        world.entity_mut(e).insert(X);
+        assert!(world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+
+        //remove `X` again and ensure that `Y` and `Z` was removed too
+        world.entity_mut(e).remove_with_requires::<X>();
+        assert!(!world.entity(e).contains::<X>());
+        assert!(!world.entity(e).contains::<Y>());
+        assert!(!world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+    }
+
+    #[test]
+    fn remove_component_and_his_required_components() {
+        #[derive(Component)]
+        #[require(Y)]
+        struct X;
+
+        #[derive(Component, Default)]
+        #[require(Z)]
+        struct Y;
+
+        #[derive(Component, Default)]
+        struct Z;
+
+        #[derive(Component)]
+        struct V;
+
+        let mut world = World::new();
+
+        let e = world.spawn((X, V)).id();
+        assert!(world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+
+        //check that `remove` works as expected
+        world.entity_mut(e).remove::<X>();
+        assert!(!world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+
+        world.entity_mut(e).insert(X);
+        assert!(world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+
+        //remove `X` again and ensure that `Y` and `Z` was removed too
+        world.entity_mut(e).remove_with_requires::<X>();
+        assert!(!world.entity(e).contains::<X>());
+        assert!(!world.entity(e).contains::<Y>());
+        assert!(!world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<V>());
+    }
+
+    #[test]
+    fn remove_bundle_and_his_required_components() {
+        #[derive(Component, Default)]
+        #[require(Y)]
+        struct X;
+
+        #[derive(Component, Default)]
+        struct Y;
+
+        #[derive(Component, Default)]
+        #[require(W)]
+        struct Z;
+
+        #[derive(Component, Default)]
+        struct W;
+
+        #[derive(Component)]
+        struct V;
+
+        #[derive(Bundle, Default)]
+        struct TestBundle {
+            x: X,
+            z: Z,
+        }
+
+        let mut world = World::new();
+        let e = world.spawn((TestBundle::default(), V)).id();
+
+        assert!(world.entity(e).contains::<X>());
+        assert!(world.entity(e).contains::<Y>());
+        assert!(world.entity(e).contains::<Z>());
+        assert!(world.entity(e).contains::<W>());
+        assert!(world.entity(e).contains::<V>());
+
+        world.entity_mut(e).remove_with_requires::<TestBundle>();
+        assert!(!world.entity(e).contains::<X>());
+        assert!(!world.entity(e).contains::<Y>());
+        assert!(!world.entity(e).contains::<Z>());
+        assert!(!world.entity(e).contains::<W>());
+        assert!(world.entity(e).contains::<V>());
+    }
+
     #[test]
     fn runtime_required_components() {
         // Same as `required_components` test but with runtime registration

crates/bevy_ecs/src/system/commands/mod.rs

@@ -1369,6 +1369,34 @@ impl EntityCommands<'_> {
         self.queue(remove::<T>)
     }
 
+    /// Removes all components in the [`Bundle`] components and remove all required components for each component in the [`Bundle`] from entity.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    ///
+    /// #[derive(Component)]
+    /// #[require(B)]
+    /// struct A;
+    /// #[derive(Component, Default)]
+    /// struct B;
+    ///
+    /// #[derive(Resource)]
+    /// struct PlayerEntity { entity: Entity }
+    ///
+    /// fn remove_with_requires_system(mut commands: Commands, player: Res<PlayerEntity>) {
+    ///     commands
+    ///         .entity(player.entity)
+    ///         // Remove both A and B components from the entity, because B is required by A
+    ///         .remove_with_requires::<A>();
+    /// }
+    /// # bevy_ecs::system::assert_is_system(remove_with_requires_system);
+    /// ```
+    pub fn remove_with_requires<T: Bundle>(&mut self) -> &mut Self {
+        self.queue(remove_with_requires::<T>)
+    }
+
     /// Removes a component from the entity.
     pub fn remove_by_id(&mut self, component_id: ComponentId) -> &mut Self {
         self.queue(remove_by_id(component_id))
@@ -1803,6 +1831,13 @@ fn remove_by_id(component_id: ComponentId) -> impl EntityCommand {
     }
 }
 
+/// An [`EntityCommand`] that remove all components in the bundle and remove all required components for each component in the bundle.
+fn remove_with_requires<T: Bundle>(entity: Entity, world: &mut World) {
+    if let Some(mut entity) = world.get_entity_mut(entity) {
+        entity.remove_with_requires::<T>();
+    }
+}
+
 /// An [`EntityCommand`] that removes all components associated with a provided entity.
 fn clear() -> impl EntityCommand {
     move |entity: Entity, world: &mut World| {
@@ -2167,6 +2202,42 @@ mod tests {
         assert!(world.contains_resource::<W<f64>>());
     }
 
+    #[test]
+    fn remove_component_with_required_components() {
+        #[derive(Component)]
+        #[require(Y)]
+        struct X;
+
+        #[derive(Component, Default)]
+        struct Y;
+
+        #[derive(Component)]
+        struct Z;
+
+        let mut world = World::default();
+        let mut queue = CommandQueue::default();
+        let e = {
+            let mut commands = Commands::new(&mut queue, &world);
+            commands.spawn((X, Z)).id()
+        };
+        queue.apply(&mut world);
+
+        assert!(world.get::<Y>(e).is_some());
+        assert!(world.get::<X>(e).is_some());
+        assert!(world.get::<Z>(e).is_some());
+
+        {
+            let mut commands = Commands::new(&mut queue, &world);
+            commands.entity(e).remove_with_requires::<X>();
+        }
+        queue.apply(&mut world);
+
+        assert!(world.get::<Y>(e).is_none());
+        assert!(world.get::<X>(e).is_none());
+
+        assert!(world.get::<Z>(e).is_some());
+    }
+
     fn is_send<T: Send>() {}
     fn is_sync<T: Sync>() {}
 

crates/bevy_ecs/src/world/entity_ref.rs

@@ -1229,6 +1229,20 @@ impl<'w> EntityWorldMut<'w> {
         self
     }
 
+    /// Removes all components in the [`Bundle`] and remove all required components for each component in the bundle
+    pub fn remove_with_requires<T: Bundle>(&mut self) -> &mut Self {
+        let storages = &mut self.world.storages;
+        let components = &mut self.world.components;
+        let bundles = &mut self.world.bundles;
+
+        let bundle_id = bundles.register_contributed_bundle_info::<T>(components, storages);
+
+        // SAFETY: the dynamic `BundleInfo` is initialized above
+        self.location = unsafe { self.remove_bundle(bundle_id) };
+
+        self
+    }
+
     /// Removes any components except those in the [`Bundle`] (and its Required Components) from the entity.
     ///
     /// See [`EntityCommands::retain`](crate::system::EntityCommands::retain) for more details.