Added modify_component to EntityWorldMut, DeferredWorld, and World (#16668)

# Objective

- Make working with immutable components more ergonomic
- Assist #16662

## Solution

Added `modify_component` to `World` and `EntityWorldMut`. This method
"removes" a component from an entity, gives a mutable reference to it to
a provided closure, and then "re-inserts" the component back onto the
entity. This replacement triggers the `OnReplace` and `OnInsert` hooks,
but does _not_ cause an archetype move, as the removal is purely
simulated.

## Testing

- Added doc-tests and a unit test.

---

## Showcase

```rust
use bevy_ecs::prelude::*;

/// An immutable component.
#[derive(Component, PartialEq, Eq, Debug)]
#[component(immutable)]
struct Foo(bool);

let mut world = World::default();

let mut entity = world.spawn(Foo(false));

assert_eq!(entity.get::<Foo>(), Some(&Foo(false)));

// Before the closure is executed, the `OnReplace` hooks/observers are triggered
entity.modify_component(|foo: &mut Foo| {
    foo.0 = true;
});
// After the closure is executed, `OnInsert` hooks/observers are triggered

assert_eq!(entity.get::<Foo>(), Some(&Foo(true)));
```

## Notes

- If the component is not available on the entity, the closure and hooks
aren't executed, and `None` is returned. I chose this as an alternative
to returning an error or panicking, but I'm open to changing that based
on feedback.
- This relies on `unsafe`, in particular for accessing the `Archetype`
to trigger hooks. All the unsafe operations are contained within
`DeferredWorld::modify_component`, and I would appreciate that this
function is given special attention to ensure soundness.
- The `OnAdd` hook can never be triggered by this method, since the
component must already be inserted. I have chosen to not trigger
`OnRemove`, as I believe it makes sense that this method is purely a
replacement operation, not an actual removal/insertion.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Malek <50841145+MalekiRe@users.noreply.github.com>

crates/bevy_ecs/src/world/deferred_world.rs

@@ -14,7 +14,7 @@ use crate::{
     world::{error::EntityFetchError, WorldEntityFetch},
 };
 
-use super::{unsafe_world_cell::UnsafeWorldCell, Mut, World};
+use super::{unsafe_world_cell::UnsafeWorldCell, Mut, World, ON_INSERT, ON_REPLACE};
 
 /// A [`World`] reference that disallows structural ECS changes.
 /// This includes initializing resources, registering components or spawning entities.
@@ -82,6 +82,88 @@ impl<'w> DeferredWorld<'w> {
         unsafe { self.world.get_entity(entity)?.get_mut() }
     }
 
+    /// Temporarily removes a [`Component`] `T` from the provided [`Entity`] and
+    /// runs the provided closure on it, returning the result if `T` was available.
+    /// This will trigger the `OnRemove` and `OnReplace` component hooks without
+    /// causing an archetype move.
+    ///
+    /// This is most useful with immutable components, where removal and reinsertion
+    /// is the only way to modify a value.
+    ///
+    /// If you do not need to ensure the above hooks are triggered, and your component
+    /// is mutable, prefer using [`get_mut`](DeferredWorld::get_mut).
+    #[inline]
+    pub(crate) fn modify_component<T: Component, R>(
+        &mut self,
+        entity: Entity,
+        f: impl FnOnce(&mut T) -> R,
+    ) -> Result<Option<R>, EntityFetchError> {
+        // If the component is not registered, then it doesn't exist on this entity, so no action required.
+        let Some(component_id) = self.component_id::<T>() else {
+            return Ok(None);
+        };
+
+        let entity_cell = match self.get_entity_mut(entity) {
+            Ok(cell) => cell,
+            Err(EntityFetchError::AliasedMutability(..)) => {
+                return Err(EntityFetchError::AliasedMutability(entity))
+            }
+            Err(EntityFetchError::NoSuchEntity(..)) => {
+                return Err(EntityFetchError::NoSuchEntity(entity, self.world))
+            }
+        };
+
+        if !entity_cell.contains::<T>() {
+            return Ok(None);
+        }
+
+        let archetype = &raw const *entity_cell.archetype();
+
+        // SAFETY:
+        // - DeferredWorld ensures archetype pointer will remain valid as no
+        //   relocations will occur.
+        // - component_id exists on this world and this entity
+        // - ON_REPLACE is able to accept ZST events
+        unsafe {
+            let archetype = &*archetype;
+            self.trigger_on_replace(archetype, entity, [component_id].into_iter());
+            if archetype.has_replace_observer() {
+                self.trigger_observers(ON_REPLACE, entity, [component_id].into_iter());
+            }
+        }
+
+        let mut entity_cell = self
+            .get_entity_mut(entity)
+            .expect("entity access confirmed above");
+
+        // SAFETY: we will run the required hooks to simulate removal/replacement.
+        let mut component = unsafe {
+            entity_cell
+                .get_mut_assume_mutable::<T>()
+                .expect("component access confirmed above")
+        };
+
+        let result = f(&mut component);
+
+        // Simulate adding this component by updating the relevant ticks
+        *component.ticks.added = *component.ticks.changed;
+
+        // SAFETY:
+        // - DeferredWorld ensures archetype pointer will remain valid as no
+        //   relocations will occur.
+        // - component_id exists on this world and this entity
+        // - ON_REPLACE is able to accept ZST events
+        unsafe {
+            let archetype = &*archetype;
+            self.trigger_on_insert(archetype, entity, [component_id].into_iter());
+            if archetype.has_insert_observer() {
+                self.trigger_observers(ON_INSERT, entity, [component_id].into_iter());
+            }
+        }
+
+        Ok(Some(result))
+    }
+
     /// Returns [`EntityMut`]s that expose read and write operations for the
     /// given `entities`, returning [`Err`] if any of the given entities do not
     /// exist. Instead of immediately unwrapping the value returned from this

crates/bevy_ecs/src/world/entity_ref.rs

@@ -1204,6 +1204,59 @@ impl<'w> EntityWorldMut<'w> {
         unsafe { self.get_mut_assume_mutable() }
     }
 
+    /// Temporarily removes a [`Component`] `T` from this [`Entity`] and runs the
+    /// provided closure on it, returning the result if `T` was available.
+    /// This will trigger the `OnRemove` and `OnReplace` component hooks without
+    /// causing an archetype move.
+    ///
+    /// This is most useful with immutable components, where removal and reinsertion
+    /// is the only way to modify a value.
+    ///
+    /// If you do not need to ensure the above hooks are triggered, and your component
+    /// is mutable, prefer using [`get_mut`](EntityWorldMut::get_mut).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// #[derive(Component, PartialEq, Eq, Debug)]
+    /// #[component(immutable)]
+    /// struct Foo(bool);
+    ///
+    /// # let mut world = World::default();
+    /// # world.register_component::<Foo>();
+    /// #
+    /// # let entity = world.spawn(Foo(false)).id();
+    /// #
+    /// # let mut entity = world.entity_mut(entity);
+    /// #
+    /// # assert_eq!(entity.get::<Foo>(), Some(&Foo(false)));
+    /// #
+    /// entity.modify_component(|foo: &mut Foo| {
+    ///     foo.0 = true;
+    /// });
+    /// #
+    /// # assert_eq!(entity.get::<Foo>(), Some(&Foo(true)));
+    /// ```
+    ///
+    /// # Panics
+    ///
+    /// If the entity has been despawned while this `EntityWorldMut` is still alive.
+    #[inline]
+    pub fn modify_component<T: Component, R>(&mut self, f: impl FnOnce(&mut T) -> R) -> Option<R> {
+        self.assert_not_despawned();
+
+        let result = self
+            .world
+            .modify_component(self.entity, f)
+            .expect("entity access must be valid")?;
+
+        self.update_location();
+
+        Some(result)
+    }
+
     /// Gets mutable access to the component of type `T` for the current entity.
     /// Returns `None` if the entity does not have a component of type `T`.
     ///
@@ -5397,4 +5450,87 @@ mod tests {
                 .unwrap()
         );
     }
+
+    #[test]
+    fn with_component_activates_hooks() {
+        use core::sync::atomic::{AtomicBool, AtomicU8, Ordering};
+
+        #[derive(Component, PartialEq, Eq, Debug)]
+        #[component(immutable)]
+        struct Foo(bool);
+
+        static EXPECTED_VALUE: AtomicBool = AtomicBool::new(false);
+
+        static ADD_COUNT: AtomicU8 = AtomicU8::new(0);
+        static REMOVE_COUNT: AtomicU8 = AtomicU8::new(0);
+        static REPLACE_COUNT: AtomicU8 = AtomicU8::new(0);
+        static INSERT_COUNT: AtomicU8 = AtomicU8::new(0);
+
+        let mut world = World::default();
+
+        world.register_component::<Foo>();
+        world
+            .register_component_hooks::<Foo>()
+            .on_add(|world, entity, _| {
+                ADD_COUNT.fetch_add(1, Ordering::Relaxed);
+
+                assert_eq!(
+                    world.get(entity),
+                    Some(&Foo(EXPECTED_VALUE.load(Ordering::Relaxed)))
+                );
+            })
+            .on_remove(|world, entity, _| {
+                REMOVE_COUNT.fetch_add(1, Ordering::Relaxed);
+
+                assert_eq!(
+                    world.get(entity),
+                    Some(&Foo(EXPECTED_VALUE.load(Ordering::Relaxed)))
+                );
+            })
+            .on_replace(|world, entity, _| {
+                REPLACE_COUNT.fetch_add(1, Ordering::Relaxed);
+
+                assert_eq!(
+                    world.get(entity),
+                    Some(&Foo(EXPECTED_VALUE.load(Ordering::Relaxed)))
+                );
+            })
+            .on_insert(|world, entity, _| {
+                INSERT_COUNT.fetch_add(1, Ordering::Relaxed);
+
+                assert_eq!(
+                    world.get(entity),
+                    Some(&Foo(EXPECTED_VALUE.load(Ordering::Relaxed)))
+                );
+            });
+
+        let entity = world.spawn(Foo(false)).id();
+
+        assert_eq!(ADD_COUNT.load(Ordering::Relaxed), 1);
+        assert_eq!(REMOVE_COUNT.load(Ordering::Relaxed), 0);
+        assert_eq!(REPLACE_COUNT.load(Ordering::Relaxed), 0);
+        assert_eq!(INSERT_COUNT.load(Ordering::Relaxed), 1);
+
+        let mut entity = world.entity_mut(entity);
+
+        let archetype_pointer_before = &raw const *entity.archetype();
+
+        assert_eq!(entity.get::<Foo>(), Some(&Foo(false)));
+
+        entity.modify_component(|foo: &mut Foo| {
+            foo.0 = true;
+            EXPECTED_VALUE.store(foo.0, Ordering::Relaxed);
+        });
+
+        let archetype_pointer_after = &raw const *entity.archetype();
+
+        assert_eq!(entity.get::<Foo>(), Some(&Foo(true)));
+
+        assert_eq!(ADD_COUNT.load(Ordering::Relaxed), 1);
+        assert_eq!(REMOVE_COUNT.load(Ordering::Relaxed), 0);
+        assert_eq!(REPLACE_COUNT.load(Ordering::Relaxed), 1);
+        assert_eq!(INSERT_COUNT.load(Ordering::Relaxed), 2);
+
+        assert_eq!(archetype_pointer_before, archetype_pointer_after);
+    }
 }

crates/bevy_ecs/src/world/mod.rs

@@ -1282,6 +1282,59 @@ impl World {
         unsafe { self.as_unsafe_world_cell().get_entity(entity)?.get_mut() }
     }
 
+    /// Temporarily removes a [`Component`] `T` from the provided [`Entity`] and
+    /// runs the provided closure on it, returning the result if `T` was available.
+    /// This will trigger the `OnRemove` and `OnReplace` component hooks without
+    /// causing an archetype move.
+    ///
+    /// This is most useful with immutable components, where removal and reinsertion
+    /// is the only way to modify a value.
+    ///
+    /// If you do not need to ensure the above hooks are triggered, and your component
+    /// is mutable, prefer using [`get_mut`](World::get_mut).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// #[derive(Component, PartialEq, Eq, Debug)]
+    /// #[component(immutable)]
+    /// struct Foo(bool);
+    ///
+    /// # let mut world = World::default();
+    /// # world.register_component::<Foo>();
+    /// #
+    /// # let entity = world.spawn(Foo(false)).id();
+    /// #
+    /// world.modify_component(entity, |foo: &mut Foo| {
+    ///     foo.0 = true;
+    /// });
+    /// #
+    /// # assert_eq!(world.get::<Foo>(entity), Some(&Foo(true)));
+    /// ```
+    #[inline]
+    pub fn modify_component<T: Component, R>(
+        &mut self,
+        entity: Entity,
+        f: impl FnOnce(&mut T) -> R,
+    ) -> Result<Option<R>, EntityFetchError> {
+        let mut world = DeferredWorld::from(&mut *self);
+
+        let result = match world.modify_component(entity, f) {
+            Ok(result) => result,
+            Err(EntityFetchError::AliasedMutability(..)) => {
+                return Err(EntityFetchError::AliasedMutability(entity))
+            }
+            Err(EntityFetchError::NoSuchEntity(..)) => {
+                return Err(EntityFetchError::NoSuchEntity(entity, self.into()))
+            }
+        };
+
+        self.flush();
+        Ok(result)
+    }
+
     /// Despawns the given `entity`, if it exists. This will also remove all of the entity's
     /// [`Component`]s. Returns `true` if the `entity` is successfully despawned and `false` if
     /// the `entity` does not exist.