Finish enhancing ReflectCommandExt to work with Bundles (bevyengine…

…#15152)

# Objective

- Finish resolving bevyengine#15125
- Inserting bundles was implemented in
bevyengine#15128 but removing bundles still
needed to be implemented.

## Solution

- Modified `bevy_ecs::reflect::entity_commands::remove_reflect` to
handle both components and bundles
- Modified documentation of `ReflectCommandExt` methods to reflect that
one can now use bundles with these commands.

## Testing

- Three tests were added to match the ones for inserting components.

crates/bevy_ecs/src/reflect/entity_commands.rs

@@ -13,16 +13,17 @@ use std::marker::PhantomData;
 
 /// An extension trait for [`EntityCommands`] for reflection related functions
 pub trait ReflectCommandExt {
-    /// Adds the given boxed reflect component to the entity using the reflection data in
+    /// Adds the given boxed reflect component or bundle to the entity using the reflection data in
     /// [`AppTypeRegistry`].
     ///
-    /// This will overwrite any previous component of the same type.
+    /// This will overwrite any previous component(s) of the same type.
     ///
     /// # Panics
     ///
     /// - If the entity doesn't exist.
-    /// - If [`AppTypeRegistry`] does not have the reflection data for the given [`Component`](crate::component::Component).
-    /// - If the component data is invalid. See [`PartialReflect::apply`] for further details.
+    /// - If [`AppTypeRegistry`] does not have the reflection data for the given
+    ///     [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle).
+    /// - If the component or bundle data is invalid. See [`PartialReflect::apply`] for further details.
     /// - If [`AppTypeRegistry`] is not present in the [`World`].
     ///
     /// # Note
@@ -38,12 +39,12 @@ pub trait ReflectCommandExt {
     /// // or write to the TypeRegistry directly to register all your components
     ///
     /// # use bevy_ecs::prelude::*;
-    /// # use bevy_ecs::reflect::ReflectCommandExt;
+    /// # use bevy_ecs::reflect::{ReflectCommandExt, ReflectBundle};
     /// # use bevy_reflect::{FromReflect, FromType, Reflect, TypeRegistry};
     /// // A resource that can hold any component that implements reflect as a boxed reflect component
     /// #[derive(Resource)]
-    /// struct Prefab{
-    ///     component: Box<dyn Reflect>,
+    /// struct Prefab {
+    ///     data: Box<dyn Reflect>,
     /// }
     /// #[derive(Component, Reflect, Default)]
     /// #[reflect(Component)]
@@ -53,23 +54,37 @@ pub trait ReflectCommandExt {
     /// #[reflect(Component)]
     /// struct ComponentB(String);
     ///
+    /// #[derive(Bundle, Reflect, Default)]
+    /// #[reflect(Bundle)]
+    /// struct BundleA {
+    ///     a: ComponentA,
+    ///     b: ComponentB,
+    /// }
+    ///
     /// fn insert_reflect_component(
     ///     mut commands: Commands,
     ///     mut prefab: ResMut<Prefab>
     ///     ) {
     ///     // Create a set of new boxed reflect components to use
     ///     let boxed_reflect_component_a: Box<dyn Reflect> = Box::new(ComponentA(916));
     ///     let boxed_reflect_component_b: Box<dyn Reflect>  = Box::new(ComponentB("NineSixteen".to_string()));
+    ///     let boxed_reflect_bundle_a: Box<dyn Reflect> = Box::new(BundleA {
+    ///         a: ComponentA(24),
+    ///         b: ComponentB("Twenty-Four".to_string()),
+    ///     });
     ///
     ///     // You can overwrite the component in the resource with either ComponentA or ComponentB
-    ///     prefab.component = boxed_reflect_component_a;
-    ///     prefab.component = boxed_reflect_component_b;
+    ///     prefab.data = boxed_reflect_component_a;
+    ///     prefab.data = boxed_reflect_component_b;
+    ///
+    ///     // Or even with BundleA
+    ///     prefab.data = boxed_reflect_bundle_a;
     ///     
-    ///     // No matter which component is in the resource and without knowing the exact type, you can
-    ///     // use the insert_reflect entity command to insert that component into an entity.
+    ///     // No matter which component or bundle is in the resource and without knowing the exact type, you can
+    ///     // use the insert_reflect entity command to insert that component/bundle into an entity.
     ///     commands
     ///         .spawn_empty()
-    ///         .insert_reflect(prefab.component.clone_value());
+    ///         .insert_reflect(prefab.data.clone_value());
     /// }
     ///
     /// ```
@@ -90,10 +105,15 @@ pub trait ReflectCommandExt {
         component: Box<dyn PartialReflect>,
     ) -> &mut Self;
 
-    /// Removes from the entity the component with the given type name registered in [`AppTypeRegistry`].
+    /// Removes from the entity the component or bundle with the given type name registered in [`AppTypeRegistry`].
+    ///
+    /// If the type is a bundle, it will remove any components in that bundle regardless if the entity
+    /// contains all the components.
     ///
-    /// Does nothing if the entity does not have a component of the same type, if [`AppTypeRegistry`]
-    /// does not contain the reflection data for the given component, or if the entity does not exist.
+    /// Does nothing if the type is a component and the entity does not have a component of the same type,
+    /// if the type is a bundle and the entity does not contain any of the components in the bundle,
+    /// if [`AppTypeRegistry`] does not contain the reflection data for the given component,
+    /// or if the entity does not exist.
     ///
     /// # Note
     ///
@@ -103,36 +123,43 @@ pub trait ReflectCommandExt {
     /// # Example
     ///
     /// ```
-    /// // Note that you need to register the component type in the AppTypeRegistry prior to using
+    /// // Note that you need to register the component/bundle type in the AppTypeRegistry prior to using
     /// // reflection. You can use the helpers on the App with `app.register_type::<ComponentA>()`
-    /// // or write to the TypeRegistry directly to register all your components
+    /// // or write to the TypeRegistry directly to register all your components and bundles
     ///
     /// # use bevy_ecs::prelude::*;
-    /// # use bevy_ecs::reflect::ReflectCommandExt;
+    /// # use bevy_ecs::reflect::{ReflectCommandExt, ReflectBundle};
     /// # use bevy_reflect::{FromReflect, FromType, Reflect, TypeRegistry};
     ///
-    /// // A resource that can hold any component that implements reflect as a boxed reflect component
+    /// // A resource that can hold any component or bundle that implements reflect as a boxed reflect
     /// #[derive(Resource)]
     /// struct Prefab{
     ///     entity: Entity,
-    ///     component: Box<dyn Reflect>,
+    ///     data: Box<dyn Reflect>,
     /// }
     /// #[derive(Component, Reflect, Default)]
     /// #[reflect(Component)]
     /// struct ComponentA(u32);
     /// #[derive(Component, Reflect, Default)]
     /// #[reflect(Component)]
     /// struct ComponentB(String);
+    /// #[derive(Bundle, Reflect, Default)]
+    /// #[reflect(Bundle)]
+    /// struct BundleA {
+    ///     a: ComponentA,
+    ///     b: ComponentB,
+    /// }
     ///
     /// fn remove_reflect_component(
     ///     mut commands: Commands,
     ///     prefab: Res<Prefab>
     ///     ) {
-    ///     // Prefab can hold any boxed reflect component. In this case either
-    ///     // ComponentA or ComponentB. No matter which component is in the resource though,
-    ///     // we can attempt to remove any component of that same type from an entity.
+    ///     // Prefab can hold any boxed reflect component or bundle. In this case either
+    ///     // ComponentA, ComponentB, or BundleA. No matter which component or bundle is in the resource though,
+    ///     // we can attempt to remove any component (or set of components in the case of a bundle)
+    ///     // of that same type from an entity.
     ///     commands.entity(prefab.entity)
-    ///         .remove_reflect(prefab.component.reflect_type_path().to_owned());
+    ///         .remove_reflect(prefab.data.reflect_type_path().to_owned());
     /// }
     ///
     /// ```
@@ -187,7 +214,7 @@ impl ReflectCommandExt for EntityCommands<'_> {
     }
 }
 
-/// Helper function to add a reflect component to a given entity
+/// Helper function to add a reflect component or bundle to a given entity
 fn insert_reflect(
     world: &mut World,
     entity: Entity,
@@ -214,14 +241,15 @@ fn insert_reflect(
     }
 }
 
-/// A [`Command`] that adds the boxed reflect component to an entity using the data in
+/// A [`Command`] that adds the boxed reflect component or bundle to an entity using the data in
 /// [`AppTypeRegistry`].
 ///
 /// See [`ReflectCommandExt::insert_reflect`] for details.
 pub struct InsertReflect {
     /// The entity on which the component will be inserted.
     pub entity: Entity,
-    /// The reflect [`Component`](crate::component::Component) that will be added to the entity.
+    /// The reflect [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle)
+    /// that will be added to the entity.
     pub component: Box<dyn PartialReflect>,
 }
 
@@ -232,7 +260,7 @@ impl Command for InsertReflect {
     }
 }
 
-/// A [`Command`] that adds the boxed reflect component to an entity using the data in the provided
+/// A [`Command`] that adds the boxed reflect component or bundle to an entity using the data in the provided
 /// [`Resource`] that implements [`AsRef<TypeRegistry>`].
 ///
 /// See [`ReflectCommandExt::insert_reflect_with_registry`] for details.
@@ -253,7 +281,7 @@ impl<T: Resource + AsRef<TypeRegistry>> Command for InsertReflectWithRegistry<T>
     }
 }
 
-/// Helper function to remove a reflect component from a given entity
+/// Helper function to remove a reflect component or bundle from a given entity
 fn remove_reflect(
     world: &mut World,
     entity: Entity,
@@ -266,20 +294,22 @@ fn remove_reflect(
     let Some(type_registration) = type_registry.get_with_type_path(&component_type_path) else {
         return;
     };
-    let Some(reflect_component) = type_registration.data::<ReflectComponent>() else {
-        return;
-    };
-    reflect_component.remove(&mut entity);
+    if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
+        reflect_component.remove(&mut entity);
+    } else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
+        reflect_bundle.remove(&mut entity);
+    }
 }
 
-/// A [`Command`] that removes the component of the same type as the given component type name from
+/// A [`Command`] that removes the component or bundle of the same type as the given type name from
 /// the provided entity.
 ///
 /// See [`ReflectCommandExt::remove_reflect`] for details.
 pub struct RemoveReflect {
     /// The entity from which the component will be removed.
     pub entity: Entity,
-    /// The [`Component`](crate::component::Component) type name that will be used to remove a component
+    /// The [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle)
+    /// type name that will be used to remove a component
     /// of the same type from the entity.
     pub component_type_path: Cow<'static, str>,
 }
@@ -296,15 +326,16 @@ impl Command for RemoveReflect {
     }
 }
 
-/// A [`Command`] that removes the component of the same type as the given component type name from
+/// A [`Command`] that removes the component or bundle of the same type as the given type name from
 /// the provided entity using the provided [`Resource`] that implements [`AsRef<TypeRegistry>`].
 ///
 /// See [`ReflectCommandExt::remove_reflect_with_registry`] for details.
 pub struct RemoveReflectWithRegistry<T: Resource + AsRef<TypeRegistry>> {
     /// The entity from which the component will be removed.
     pub entity: Entity,
     pub _t: PhantomData<T>,
-    /// The [`Component`](crate::component::Component) type name that will be used to remove a component
+    /// The [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle)
+    /// type name that will be used to remove a component
     /// of the same type from the entity.
     pub component_type_name: Cow<'static, str>,
 }
@@ -505,4 +536,112 @@ mod tests {
         assert_eq!(world.get::<ComponentA>(entity), Some(&ComponentA(31)));
         assert_eq!(world.get::<ComponentB>(entity), Some(&ComponentB(20)));
     }
+
+    #[test]
+    fn insert_reflect_bundle_with_registry() {
+        let mut world = World::new();
+
+        let mut type_registry = TypeRegistryResource {
+            type_registry: TypeRegistry::new(),
+        };
+
+        type_registry.type_registry.register::<BundleA>();
+        type_registry
+            .type_registry
+            .register_type_data::<BundleA, ReflectBundle>();
+        world.insert_resource(type_registry);
+
+        let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
+        let mut commands = system_state.get_mut(&mut world);
+
+        let entity = commands.spawn_empty().id();
+        let bundle = Box::new(BundleA {
+            a: ComponentA(31),
+            b: ComponentB(20),
+        }) as Box<dyn PartialReflect>;
+
+        commands
+            .entity(entity)
+            .insert_reflect_with_registry::<TypeRegistryResource>(bundle);
+        system_state.apply(&mut world);
+
+        assert_eq!(world.get::<ComponentA>(entity), Some(&ComponentA(31)));
+        assert_eq!(world.get::<ComponentB>(entity), Some(&ComponentB(20)));
+    }
+
+    #[test]
+    fn remove_reflected_bundle() {
+        let mut world = World::new();
+
+        let type_registry = AppTypeRegistry::default();
+        {
+            let mut registry = type_registry.write();
+            registry.register::<BundleA>();
+            registry.register_type_data::<BundleA, ReflectBundle>();
+        }
+        world.insert_resource(type_registry);
+
+        let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
+        let mut commands = system_state.get_mut(&mut world);
+
+        let entity = commands
+            .spawn(BundleA {
+                a: ComponentA(31),
+                b: ComponentB(20),
+            })
+            .id();
+
+        let boxed_reflect_bundle_a = Box::new(BundleA {
+            a: ComponentA(1),
+            b: ComponentB(23),
+        }) as Box<dyn Reflect>;
+
+        commands
+            .entity(entity)
+            .remove_reflect(boxed_reflect_bundle_a.reflect_type_path().to_owned());
+        system_state.apply(&mut world);
+
+        assert_eq!(world.entity(entity).get::<ComponentA>(), None);
+        assert_eq!(world.entity(entity).get::<ComponentB>(), None);
+    }
+
+    #[test]
+    fn remove_reflected_bundle_with_registry() {
+        let mut world = World::new();
+
+        let mut type_registry = TypeRegistryResource {
+            type_registry: TypeRegistry::new(),
+        };
+
+        type_registry.type_registry.register::<BundleA>();
+        type_registry
+            .type_registry
+            .register_type_data::<BundleA, ReflectBundle>();
+        world.insert_resource(type_registry);
+
+        let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
+        let mut commands = system_state.get_mut(&mut world);
+
+        let entity = commands
+            .spawn(BundleA {
+                a: ComponentA(31),
+                b: ComponentB(20),
+            })
+            .id();
+
+        let boxed_reflect_bundle_a = Box::new(BundleA {
+            a: ComponentA(1),
+            b: ComponentB(23),
+        }) as Box<dyn Reflect>;
+
+        commands
+            .entity(entity)
+            .remove_reflect_with_registry::<TypeRegistryResource>(
+                boxed_reflect_bundle_a.reflect_type_path().to_owned(),
+            );
+        system_state.apply(&mut world);
+
+        assert_eq!(world.entity(entity).get::<ComponentA>(), None);
+        assert_eq!(world.entity(entity).get::<ComponentB>(), None);
+    }
 }