Make some ReflectComponent/ReflectBundle methods work with EntityMut too (#12895)

# Objective

- Make `ReflectComponent::apply`, `ReflectComponent::reflect_mut` and
`ReflectBundle::apply` work with `EntityMut` too (currently they only
work with the more restricting `EntityWorldMut`);
- Note: support for the `Filtered*` variants has been left out since the
conversion in that case is more expensive. Let me know if I should add
support for them too.

## Solution

- Make `ReflectComponent::apply`, `ReflectComponent::reflect_mut` and
`ReflectBundle::apply` take an `impl Into<EntityMut<'a>>`;
- Make the corresponding `*Fns` function pointers take a `EntityMut`.

---

## Changelog

- `ReflectComponent::apply`, `ReflectComponent::reflect_mut` and
`ReflectBundle::apply` now accept `EntityMut` as well

## Migration Guide

- `ReflectComponentFns`'s `apply` and `reflect_mut` fields now take
`EntityMut` instead of `&mut EntityWorldMut`
- `ReflectBundleFns`'s `apply` field now takes `EntityMut` instead of
`&mut EntityWorldMut`

crates/bevy_ecs/src/reflect/bundle.rs

@@ -6,7 +6,10 @@
 //! Same as [`super::component`], but for bundles.
 use std::any::TypeId;
 
-use crate::{prelude::Bundle, world::EntityWorldMut};
+use crate::{
+    prelude::Bundle,
+    world::{EntityMut, EntityWorldMut},
+};
 use bevy_reflect::{FromReflect, FromType, Reflect, ReflectRef, TypeRegistry};
 
 use super::ReflectComponent;
@@ -26,7 +29,7 @@ pub struct ReflectBundleFns {
     /// Function pointer implementing [`ReflectBundle::insert()`].
     pub insert: fn(&mut EntityWorldMut, &dyn Reflect),
     /// Function pointer implementing [`ReflectBundle::apply()`].
-    pub apply: fn(&mut EntityWorldMut, &dyn Reflect, &TypeRegistry),
+    pub apply: fn(EntityMut, &dyn Reflect, &TypeRegistry),
     /// Function pointer implementing [`ReflectBundle::apply_or_insert()`].
     pub apply_or_insert: fn(&mut EntityWorldMut, &dyn Reflect, &TypeRegistry),
     /// Function pointer implementing [`ReflectBundle::remove()`].
@@ -55,13 +58,13 @@ impl ReflectBundle {
     /// # Panics
     ///
     /// Panics if there is no [`Bundle`] of the given type.
-    pub fn apply(
+    pub fn apply<'a>(
         &self,
-        entity: &mut EntityWorldMut,
+        entity: impl Into<EntityMut<'a>>,
         bundle: &dyn Reflect,
         registry: &TypeRegistry,
     ) {
-        (self.0.apply)(entity, bundle, registry);
+        (self.0.apply)(entity.into(), bundle, registry);
     }
 
     /// Uses reflection to set the value of this [`Bundle`] type in the entity to the given value or insert a new one if it does not exist.
@@ -121,7 +124,7 @@ impl<B: Bundle + Reflect + FromReflect> FromType<B> for ReflectBundle {
                 let bundle = B::from_reflect(reflected_bundle).unwrap();
                 entity.insert(bundle);
             },
-            apply: |entity, reflected_bundle, registry| {
+            apply: |mut entity, reflected_bundle, registry| {
                 if let Some(reflect_component) =
                     registry.get_type_data::<ReflectComponent>(TypeId::of::<B>())
                 {
@@ -130,10 +133,10 @@ impl<B: Bundle + Reflect + FromReflect> FromType<B> for ReflectBundle {
                     match reflected_bundle.reflect_ref() {
                         ReflectRef::Struct(bundle) => bundle
                             .iter_fields()
-                            .for_each(|field| insert_field(entity, field, registry)),
+                            .for_each(|field| apply_field(&mut entity, field, registry)),
                         ReflectRef::Tuple(bundle) => bundle
                             .iter_fields()
-                            .for_each(|field| insert_field(entity, field, registry)),
+                            .for_each(|field| apply_field(&mut entity, field, registry)),
                         _ => panic!(
                             "expected bundle `{}` to be named struct or tuple",
                             // FIXME: once we have unique reflect, use `TypePath`.
@@ -170,29 +173,16 @@ impl<B: Bundle + Reflect + FromReflect> FromType<B> for ReflectBundle {
     }
 }
 
-fn insert_field(entity: &mut EntityWorldMut, field: &dyn Reflect, registry: &TypeRegistry) {
+fn apply_field(entity: &mut EntityMut, field: &dyn Reflect, registry: &TypeRegistry) {
     if let Some(reflect_component) = registry.get_type_data::<ReflectComponent>(field.type_id()) {
-        reflect_component.apply(entity, field);
+        reflect_component.apply(entity.reborrow(), field);
     } else if let Some(reflect_bundle) = registry.get_type_data::<ReflectBundle>(field.type_id()) {
-        reflect_bundle.apply(entity, field, registry);
+        reflect_bundle.apply(entity.reborrow(), field, registry);
     } else {
-        let is_component = entity
+        panic!(
-            .world()
+            "no `ReflectComponent` nor `ReflectBundle` registration found for `{}`",
-            .components()
+            field.reflect_type_path()
-            .get_id(field.type_id())
+        );
-            .is_some();
-
-        if is_component {
-            panic!(
-                "no `ReflectComponent` registration found for `{}`",
-                field.reflect_type_path(),
-            );
-        } else {
-            panic!(
-                "no `ReflectBundle` registration found for `{}`",
-                field.reflect_type_path(),
-            )
-        }
     }
 }
 

crates/bevy_ecs/src/reflect/component.rs

@@ -62,7 +62,7 @@ use crate::{
     change_detection::Mut,
     component::Component,
     entity::Entity,
-    world::{unsafe_world_cell::UnsafeEntityCell, EntityRef, EntityWorldMut, World},
+    world::{unsafe_world_cell::UnsafeEntityCell, EntityMut, EntityRef, EntityWorldMut, World},
 };
 use bevy_reflect::{FromReflect, FromType, Reflect, TypeRegistry};
 
@@ -98,7 +98,7 @@ pub struct ReflectComponentFns {
     /// Function pointer implementing [`ReflectComponent::insert()`].
     pub insert: fn(&mut EntityWorldMut, &dyn Reflect, &TypeRegistry),
     /// Function pointer implementing [`ReflectComponent::apply()`].
-    pub apply: fn(&mut EntityWorldMut, &dyn Reflect),
+    pub apply: fn(EntityMut, &dyn Reflect),
     /// Function pointer implementing [`ReflectComponent::apply_or_insert()`].
     pub apply_or_insert: fn(&mut EntityWorldMut, &dyn Reflect, &TypeRegistry),
     /// Function pointer implementing [`ReflectComponent::remove()`].
@@ -108,7 +108,7 @@ pub struct ReflectComponentFns {
     /// Function pointer implementing [`ReflectComponent::reflect()`].
     pub reflect: fn(EntityRef) -> Option<&dyn Reflect>,
     /// Function pointer implementing [`ReflectComponent::reflect_mut()`].
-    pub reflect_mut: for<'a> fn(&'a mut EntityWorldMut<'_>) -> Option<Mut<'a, dyn Reflect>>,
+    pub reflect_mut: fn(EntityMut) -> Option<Mut<dyn Reflect>>,
     /// Function pointer implementing [`ReflectComponent::reflect_unchecked_mut()`].
     ///
     /// # Safety
@@ -145,8 +145,8 @@ impl ReflectComponent {
     /// # Panics
     ///
     /// Panics if there is no [`Component`] of the given type.
-    pub fn apply(&self, entity: &mut EntityWorldMut, component: &dyn Reflect) {
+    pub fn apply<'a>(&self, entity: impl Into<EntityMut<'a>>, component: &dyn Reflect) {
-        (self.0.apply)(entity, component);
+        (self.0.apply)(entity.into(), component);
     }
 
     /// Uses reflection to set the value of this [`Component`] type in the entity to the given value or insert a new one if it does not exist.
@@ -177,9 +177,9 @@ impl ReflectComponent {
     /// Gets the value of this [`Component`] type from the entity as a mutable reflected reference.
     pub fn reflect_mut<'a>(
         &self,
-        entity: &'a mut EntityWorldMut<'_>,
+        entity: impl Into<EntityMut<'a>>,
     ) -> Option<Mut<'a, dyn Reflect>> {
-        (self.0.reflect_mut)(entity)
+        (self.0.reflect_mut)(entity.into())
     }
 
     /// # Safety
@@ -262,7 +262,7 @@ impl<C: Component + Reflect + FromReflect> FromType<C> for ReflectComponent {
                 });
                 entity.insert(component);
             },
-            apply: |entity, reflected_component| {
+            apply: |mut entity, reflected_component| {
                 let mut component = entity.get_mut::<C>().unwrap();
                 component.apply(reflected_component);
             },
@@ -290,7 +290,7 @@ impl<C: Component + Reflect + FromReflect> FromType<C> for ReflectComponent {
             },
             reflect: |entity| entity.get::<C>().map(|c| c as &dyn Reflect),
             reflect_mut: |entity| {
-                entity.get_mut::<C>().map(|c| Mut {
+                entity.into_mut::<C>().map(|c| Mut {
                     value: c.value as &mut dyn Reflect,
                     ticks: c.ticks,
                 })