Add EntityWorldMut::(try_)resource_scope (#20162)

# Objective

Fixes #20139

## Solution

Implement the methods, and leverage them where applicable

## Testing

Added unit tests

---

## Showcase

```rust
let id = entity_world_mut.id();
let world = entity_world_mut.into_world_mut();
world.resource_scope::<_, _>(|world, res| {
    let entity_world_mut = world.entity_mut(id);
    /* ... */
});
```

becomes

```rust
entity_world_mut.resource_scope::<_, _>(|entity, res| {
    /* ... */
});
```

Author: BigWingBeat

crates/bevy_ecs/src/event/update.rs

@@ -31,20 +31,16 @@ pub fn signal_event_update_system(signal: Option<ResMut<EventRegistry>>) {
 
 /// A system that calls [`Events::update`](super::Events::update) on all registered [`Events`][super::Events] in the world.
 pub fn event_update_system(world: &mut World, mut last_change_tick: Local<Tick>) {
-    if world.contains_resource::<EventRegistry>() {
-        world.resource_scope(|world, mut registry: Mut<EventRegistry>| {
-            registry.run_updates(world, *last_change_tick);
+    world.try_resource_scope(|world, mut registry: Mut<EventRegistry>| {
+        registry.run_updates(world, *last_change_tick);
 
-            registry.should_update = match registry.should_update {
-                // If we're always updating, keep doing so.
-                ShouldUpdateEvents::Always => ShouldUpdateEvents::Always,
-                // Disable the system until signal_event_update_system runs again.
-                ShouldUpdateEvents::Waiting | ShouldUpdateEvents::Ready => {
-                    ShouldUpdateEvents::Waiting
-                }
-            };
-        });
-    }
+        registry.should_update = match registry.should_update {
+            // If we're always updating, keep doing so.
+            ShouldUpdateEvents::Always => ShouldUpdateEvents::Always,
+            // Disable the system until signal_event_update_system runs again.
+            ShouldUpdateEvents::Waiting | ShouldUpdateEvents::Ready => ShouldUpdateEvents::Waiting,
+        };
+    });
     *last_change_tick = world.change_tick();
 }
 

crates/bevy_ecs/src/reflect/entity_commands.rs

@@ -1,10 +1,9 @@
 use crate::{
-    entity::Entity,
     prelude::Mut,
     reflect::{AppTypeRegistry, ReflectBundle, ReflectComponent},
     resource::Resource,
     system::EntityCommands,
-    world::{EntityWorldMut, World},
+    world::EntityWorldMut,
 };
 use alloc::{borrow::Cow, boxed::Box};
 use bevy_reflect::{PartialReflect, TypeRegistry};
@@ -22,7 +21,7 @@ pub trait ReflectCommandExt {
     /// - 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`].
+    /// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
     ///
     /// # Note
     ///
@@ -92,11 +91,11 @@ pub trait ReflectCommandExt {
     ///
     /// # Panics
     ///
-    /// - If the given [`Resource`] is not present in the [`World`].
+    /// - If the given [`Resource`] is not present in the [`World`](crate::world::World).
     ///
     /// # Note
     ///
-    /// - The given [`Resource`] is removed from the [`World`] before the command is applied.
+    /// - The given [`Resource`] is removed from the [`World`](crate::world::World) before the command is applied.
     fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
         &mut self,
         component: Box<dyn PartialReflect>,
@@ -214,23 +213,18 @@ impl<'w> EntityWorldMut<'w> {
     /// - 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`].
+    /// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
     ///
     /// # Note
     ///
     /// Prefer to use the typed [`EntityWorldMut::insert`] if possible. Adding a reflected component
     /// is much slower.
     pub fn insert_reflect(&mut self, component: Box<dyn PartialReflect>) -> &mut Self {
         self.assert_not_despawned();
-        let entity_id = self.id();
-        self.world_scope(|world| {
-            world.resource_scope(|world, registry: Mut<AppTypeRegistry>| {
-                let type_registry = &registry.as_ref().read();
-                insert_reflect_with_registry_ref(world, entity_id, type_registry, component);
-            });
-            world.flush();
+        self.resource_scope(|entity, registry: Mut<AppTypeRegistry>| {
+            let type_registry = &registry.as_ref().read();
+            insert_reflect_with_registry_ref(entity, type_registry, component);
         });
-        self.update_location();
         self
     }
 
@@ -245,21 +239,16 @@ impl<'w> EntityWorldMut<'w> {
     /// - If the given [`Resource`] 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 the given [`Resource`] is not present in the [`World`].
+    /// - If the given [`Resource`] is not present in the [`World`](crate::world::World).
     pub fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
         &mut self,
         component: Box<dyn PartialReflect>,
     ) -> &mut Self {
         self.assert_not_despawned();
-        let entity_id = self.id();
-        self.world_scope(|world| {
-            world.resource_scope(|world, registry: Mut<T>| {
-                let type_registry = registry.as_ref().as_ref();
-                insert_reflect_with_registry_ref(world, entity_id, type_registry, component);
-            });
-            world.flush();
+        self.resource_scope(|entity, registry: Mut<T>| {
+            let type_registry = registry.as_ref().as_ref();
+            insert_reflect_with_registry_ref(entity, type_registry, component);
         });
-        self.update_location();
         self
     }
 
@@ -275,28 +264,18 @@ impl<'w> EntityWorldMut<'w> {
     /// # Panics
     ///
     /// - If the entity has been despawned while this `EntityWorldMut` is still alive.
-    /// - If [`AppTypeRegistry`] is not present in the [`World`].
+    /// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
     ///
     /// # Note
     ///
     /// Prefer to use the typed [`EntityCommands::remove`] if possible. Removing a reflected component
     /// is much slower.
     pub fn remove_reflect(&mut self, component_type_path: Cow<'static, str>) -> &mut Self {
         self.assert_not_despawned();
-        let entity_id = self.id();
-        self.world_scope(|world| {
-            world.resource_scope(|world, registry: Mut<AppTypeRegistry>| {
-                let type_registry = &registry.as_ref().read();
-                remove_reflect_with_registry_ref(
-                    world,
-                    entity_id,
-                    type_registry,
-                    component_type_path,
-                );
-            });
-            world.flush();
+        self.resource_scope(|entity, registry: Mut<AppTypeRegistry>| {
+            let type_registry = &registry.as_ref().read();
+            remove_reflect_with_registry_ref(entity, type_registry, component_type_path);
         });
-        self.update_location();
         self
     }
 
@@ -313,75 +292,56 @@ impl<'w> EntityWorldMut<'w> {
     /// # Panics
     ///
     /// - If the entity has been despawned while this `EntityWorldMut` is still alive.
-    /// - If [`AppTypeRegistry`] is not present in the [`World`].
+    /// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
     pub fn remove_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
         &mut self,
         component_type_path: Cow<'static, str>,
     ) -> &mut Self {
         self.assert_not_despawned();
-        let entity_id = self.id();
-        self.world_scope(|world| {
-            world.resource_scope(|world, registry: Mut<T>| {
-                let type_registry = registry.as_ref().as_ref();
-                remove_reflect_with_registry_ref(
-                    world,
-                    entity_id,
-                    type_registry,
-                    component_type_path,
-                );
-            });
-            world.flush();
+        self.resource_scope(|entity, registry: Mut<T>| {
+            let type_registry = registry.as_ref().as_ref();
+            remove_reflect_with_registry_ref(entity, type_registry, component_type_path);
         });
-        self.update_location();
         self
     }
 }
 
 /// Helper function to add a reflect component or bundle to a given entity
 fn insert_reflect_with_registry_ref(
-    world: &mut World,
-    entity: Entity,
+    entity: &mut EntityWorldMut,
     type_registry: &TypeRegistry,
     component: Box<dyn PartialReflect>,
 ) {
     let type_info = component
         .get_represented_type_info()
         .expect("component should represent a type.");
     let type_path = type_info.type_path();
-    let Ok(mut entity) = world.get_entity_mut(entity) else {
-        panic!("error[B0003]: Could not insert a reflected component (of type {type_path}) for entity {entity}, which {}. See: https://bevy.org/learn/errors/b0003",
-        world.entities().entity_does_not_exist_error_details(entity));
-    };
     let Some(type_registration) = type_registry.get(type_info.type_id()) else {
         panic!("`{type_path}` should be registered in type registry via `App::register_type<{type_path}>`");
     };
 
     if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
-        reflect_component.insert(&mut entity, component.as_partial_reflect(), type_registry);
+        reflect_component.insert(entity, component.as_partial_reflect(), type_registry);
     } else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
-        reflect_bundle.insert(&mut entity, component.as_partial_reflect(), type_registry);
+        reflect_bundle.insert(entity, component.as_partial_reflect(), type_registry);
     } else {
         panic!("`{type_path}` should have #[reflect(Component)] or #[reflect(Bundle)]");
     }
 }
 
 /// Helper function to remove a reflect component or bundle from a given entity
 fn remove_reflect_with_registry_ref(
-    world: &mut World,
-    entity: Entity,
+    entity: &mut EntityWorldMut,
     type_registry: &TypeRegistry,
     component_type_path: Cow<'static, str>,
 ) {
-    let Ok(mut entity) = world.get_entity_mut(entity) else {
-        return;
-    };
     let Some(type_registration) = type_registry.get_with_type_path(&component_type_path) else {
         return;
     };
     if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
-        reflect_component.remove(&mut entity);
+        reflect_component.remove(entity);
     } else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
-        reflect_bundle.remove(&mut entity);
+        reflect_bundle.remove(entity);
     }
 }
 

crates/bevy_ecs/src/world/entity_ref.rs

@@ -1547,6 +1547,50 @@ impl<'w> EntityWorldMut<'w> {
         self.world.get_resource_mut()
     }
 
+    /// Temporarily removes the requested resource from the [`World`], runs custom user code,
+    /// then re-adds the resource before returning.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the resource does not exist.
+    /// Use [`try_resource_scope`](Self::try_resource_scope) instead if you want to handle this case.
+    ///
+    /// See [`World::resource_scope`] for further details.
+    #[track_caller]
+    pub fn resource_scope<R: Resource, U>(
+        &mut self,
+        f: impl FnOnce(&mut EntityWorldMut, Mut<R>) -> U,
+    ) -> U {
+        let id = self.id();
+        self.world_scope(|world| {
+            world.resource_scope(|world, res| {
+                // Acquiring a new EntityWorldMut here and using that instead of `self` is fine because
+                // the outer `world_scope` will handle updating our location if it gets changed by the user code
+                let mut this = world.entity_mut(id);
+                f(&mut this, res)
+            })
+        })
+    }
+
+    /// Temporarily removes the requested resource from the [`World`] if it exists, runs custom user code,
+    /// then re-adds the resource before returning. Returns `None` if the resource does not exist in the [`World`].
+    ///
+    /// See [`World::try_resource_scope`] for further details.
+    pub fn try_resource_scope<R: Resource, U>(
+        &mut self,
+        f: impl FnOnce(&mut EntityWorldMut, Mut<R>) -> U,
+    ) -> Option<U> {
+        let id = self.id();
+        self.world_scope(|world| {
+            world.try_resource_scope(|world, res| {
+                // Acquiring a new EntityWorldMut here and using that instead of `self` is fine because
+                // the outer `world_scope` will handle updating our location if it gets changed by the user code
+                let mut this = world.entity_mut(id);
+                f(&mut this, res)
+            })
+        })
+    }
+
     /// Retrieves the change ticks for the given component. This can be useful for implementing change
     /// detection in custom runtimes.
     ///
@@ -4959,6 +5003,46 @@ mod tests {
         assert!(entity.get_mut_by_id(invalid_component_id).is_err());
     }
 
+    #[derive(Resource)]
+    struct R(usize);
+
+    #[test]
+    fn entity_mut_resource_scope() {
+        // Keep in sync with the `resource_scope` test in lib.rs
+        let mut world = World::new();
+        let mut entity = world.spawn_empty();
+
+        assert!(entity.try_resource_scope::<R, _>(|_, _| {}).is_none());
+        entity.world_scope(|world| world.insert_resource(R(0)));
+        entity.resource_scope(|entity: &mut EntityWorldMut, mut value: Mut<R>| {
+            value.0 += 1;
+            assert!(!entity.world().contains_resource::<R>());
+        });
+        assert_eq!(entity.resource::<R>().0, 1);
+    }
+
+    #[test]
+    fn entity_mut_resource_scope_panic() {
+        let mut world = World::new();
+        world.insert_resource(R(0));
+
+        let mut entity = world.spawn_empty();
+        let old_location = entity.location();
+        let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
+            entity.resource_scope(|entity: &mut EntityWorldMut, _: Mut<R>| {
+                // Change the entity's `EntityLocation`.
+                entity.insert(TestComponent(0));
+
+                // Ensure that the entity location still gets updated even in case of a panic.
+                panic!("this should get caught by the outer scope")
+            });
+        }));
+        assert!(result.is_err());
+
+        // Ensure that the location has been properly updated.
+        assert_ne!(entity.location(), old_location);
+    }
+
     // regression test for https://github.com/bevyengine/bevy/pull/7387
     #[test]
     fn entity_mut_world_scope_panic() {
@@ -4983,6 +5067,28 @@ mod tests {
         assert_ne!(entity.location(), old_location);
     }
 
+    #[test]
+    fn entity_mut_reborrow_scope_panic() {
+        let mut world = World::new();
+
+        let mut entity = world.spawn_empty();
+        let old_location = entity.location();
+        let res = std::panic::catch_unwind(AssertUnwindSafe(|| {
+            entity.reborrow_scope(|mut entity| {
+                // Change the entity's `EntityLocation`, which invalidates the original `EntityWorldMut`.
+                // This will get updated at the end of the scope.
+                entity.insert(TestComponent(0));
+
+                // Ensure that the entity location still gets updated even in case of a panic.
+                panic!("this should get caught by the outer scope")
+            });
+        }));
+        assert!(res.is_err());
+
+        // Ensure that the location has been properly updated.
+        assert_ne!(entity.location(), old_location);
+    }
+
     // regression test for https://github.com/bevyengine/bevy/pull/7805
     #[test]
     fn removing_sparse_updates_archetype_row() {
@@ -5420,9 +5526,6 @@ mod tests {
     #[derive(Component)]
     struct A;
 
-    #[derive(Resource)]
-    struct R;
-
     #[test]
     fn disjoint_access() {
         fn disjoint_readonly(_: Query<EntityMut, With<A>>, _: Query<EntityRef, Without<A>>) {}

crates/bevy_ecs/src/world/mod.rs

@@ -2551,6 +2551,11 @@ impl World {
     /// This enables safe simultaneous mutable access to both a resource and the rest of the [`World`].
     /// For more complex access patterns, consider using [`SystemState`](crate::system::SystemState).
     ///
+    /// # Panics
+    ///
+    /// Panics if the resource does not exist.
+    /// Use [`try_resource_scope`](Self::try_resource_scope) instead if you want to handle this case.
+    ///
     /// # Example
     /// ```
     /// use bevy_ecs::prelude::*;
@@ -2568,8 +2573,6 @@ impl World {
     /// });
     /// assert_eq!(world.get_resource::<A>().unwrap().0, 2);
     /// ```
-    ///
-    /// See also [`try_resource_scope`](Self::try_resource_scope).
     #[track_caller]
     pub fn resource_scope<R: Resource, U>(&mut self, f: impl FnOnce(&mut World, Mut<R>) -> U) -> U {
         self.try_resource_scope(f)