Add methods to work with dynamic immutable components (#18532)

# Objective

- Fixes #16861

## Solution

- Added: 
  - `UnsafeEntityCell::get_mut_assume_mutable_by_id`
  - `EntityMut::get_mut_assume_mutable_by_id`
  - `EntityMut::get_mut_assume_mutable_by_id_unchecked`
  - `EntityWorldMut::into_mut_assume_mutable_by_id`
  - `EntityWorldMut::into_mut_assume_mutable`
  - `EntityWorldMut::get_mut_assume_mutable_by_id`
  - `EntityWorldMut::into_mut_assume_mutable_by_id`
  - `EntityWorldMut::modify_component_by_id`
  - `World::modify_component_by_id`
  - `DeferredWorld::modify_component_by_id`
- Added `fetch_mut_assume_mutable` to `DynamicComponentFetch` trait
(this is a breaking change)

## Testing

- CI

---

## Migration Guide

If you had previously implemented `DynamicComponentFetch` you must now
include a definition for `fetch_mut_assume_mutable`. In general this
will be identical to `fetch_mut` using the relevant alternatives for
actually getting a component.

---

## Notes

All of the added methods are minor variations on existing functions and
should therefore be of low risk for inclusion during the RC process.

crates/bevy_ecs/src/world/deferred_world.rs

@@ -103,9 +103,38 @@ impl<'w> DeferredWorld<'w> {
             return Ok(None);
         };
 
+        self.modify_component_by_id(entity, component_id, move |component| {
+            // SAFETY: component matches the component_id collected in the above line
+            let mut component = unsafe { component.with_type::<T>() };
+
+            f(&mut component)
+        })
+    }
+
+    /// Temporarily removes a [`Component`] identified by the provided
+    /// [`ComponentId`] from the provided [`Entity`] and runs the provided
+    /// closure on it, returning the result if the component 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_by_id`](DeferredWorld::get_mut_by_id).
+    ///
+    /// You should prefer the typed [`modify_component`](DeferredWorld::modify_component)
+    /// whenever possible.
+    #[inline]
+    pub(crate) fn modify_component_by_id<R>(
+        &mut self,
+        entity: Entity,
+        component_id: ComponentId,
+        f: impl for<'a> FnOnce(MutUntyped<'a>) -> R,
+    ) -> Result<Option<R>, EntityMutableFetchError> {
         let entity_cell = self.get_entity_mut(entity)?;
 
-        if !entity_cell.contains::<T>() {
+        if !entity_cell.contains_id(component_id) {
             return Ok(None);
         }
 
@@ -142,11 +171,11 @@ impl<'w> DeferredWorld<'w> {
         // SAFETY: we will run the required hooks to simulate removal/replacement.
         let mut component = unsafe {
             entity_cell
-                .get_mut_assume_mutable::<T>()
+                .get_mut_assume_mutable_by_id(component_id)
                 .expect("component access confirmed above")
         };
 
-        let result = f(&mut component);
+        let result = f(component.reborrow());
 
         // Simulate adding this component by updating the relevant ticks
         *component.ticks.added = *component.ticks.changed;

crates/bevy_ecs/src/world/entity_ref.rs

@@ -824,6 +824,39 @@ impl<'w> EntityMut<'w> {
         unsafe { component_ids.fetch_mut(self.cell) }
     }
 
+    /// Returns [untyped mutable reference(s)](MutUntyped) to component(s) for
+    /// the current entity, based on the given [`ComponentId`]s.
+    /// Assumes the given [`ComponentId`]s refer to mutable components.
+    ///
+    /// **You should prefer to use the typed API [`EntityMut::get_mut_assume_mutable`] where
+    /// possible and only use this in cases where the actual component types
+    /// are not known at compile time.**
+    ///
+    /// Unlike [`EntityMut::get_mut_assume_mutable`], this returns untyped reference(s) to
+    /// component(s), and it's the job of the caller to ensure the correct
+    /// type(s) are dereferenced (if necessary).
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if the entity does
+    ///   not have a component.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component
+    ///   is requested multiple times.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the provided [`ComponentId`]s must refer to mutable components.
+    #[inline]
+    pub unsafe fn get_mut_assume_mutable_by_id<F: DynamicComponentFetch>(
+        &mut self,
+        component_ids: F,
+    ) -> Result<F::Mut<'_>, EntityComponentError> {
+        // SAFETY:
+        // - `&mut self` ensures that no references exist to this entity's components.
+        // - We have exclusive access to all components of this entity.
+        unsafe { component_ids.fetch_mut_assume_mutable(self.cell) }
+    }
+
     /// Returns [untyped mutable reference](MutUntyped) to component for
     /// the current entity, based on the given [`ComponentId`].
     ///
@@ -852,6 +885,36 @@ impl<'w> EntityMut<'w> {
         unsafe { component_ids.fetch_mut(self.cell) }
     }
 
+    /// Returns [untyped mutable reference](MutUntyped) to component for
+    /// the current entity, based on the given [`ComponentId`].
+    /// Assumes the given [`ComponentId`]s refer to mutable components.
+    ///
+    /// Unlike [`EntityMut::get_mut_assume_mutable_by_id`], this method borrows &self instead of
+    /// &mut self, allowing the caller to access multiple components simultaneously.
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if the entity does
+    ///   not have a component.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component
+    ///   is requested multiple times.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeEntityCell`] has permission to access the component mutably
+    /// - no other references to the component exist at the same time
+    /// - the provided [`ComponentId`]s must refer to mutable components.
+    #[inline]
+    pub unsafe fn get_mut_assume_mutable_by_id_unchecked<F: DynamicComponentFetch>(
+        &self,
+        component_ids: F,
+    ) -> Result<F::Mut<'_>, EntityComponentError> {
+        // SAFETY:
+        // - The caller must ensure simultaneous access is limited
+        // - to components that are mutually independent.
+        unsafe { component_ids.fetch_mut_assume_mutable(self.cell) }
+    }
+
     /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped)
     /// to component(s) with lifetime `'w` for the current entity, based on the
     /// given [`ComponentId`]s.
@@ -885,6 +948,40 @@ impl<'w> EntityMut<'w> {
         unsafe { component_ids.fetch_mut(self.cell) }
     }
 
+    /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped)
+    /// to component(s) with lifetime `'w` for the current entity, based on the
+    /// given [`ComponentId`]s.
+    /// Assumes the given [`ComponentId`]s refer to mutable components.
+    ///
+    /// **You should prefer to use the typed API [`EntityMut::into_mut_assume_mutable`] where
+    /// possible and only use this in cases where the actual component types
+    /// are not known at compile time.**
+    ///
+    /// Unlike [`EntityMut::into_mut_assume_mutable`], this returns untyped reference(s) to
+    /// component(s), and it's the job of the caller to ensure the correct
+    /// type(s) are dereferenced (if necessary).
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if the entity does
+    ///   not have a component.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component
+    ///   is requested multiple times.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the provided [`ComponentId`]s must refer to mutable components.
+    #[inline]
+    pub unsafe fn into_mut_assume_mutable_by_id<F: DynamicComponentFetch>(
+        self,
+        component_ids: F,
+    ) -> Result<F::Mut<'w>, EntityComponentError> {
+        // SAFETY:
+        // - consuming `self` ensures that no references exist to this entity's components.
+        // - We have exclusive access to all components of this entity.
+        unsafe { component_ids.fetch_mut_assume_mutable(self.cell) }
+    }
+
     /// Returns the source code location from which this entity has been spawned.
     pub fn spawned_by(&self) -> MaybeLocation {
         self.cell.spawned_by()
@@ -1302,6 +1399,38 @@ impl<'w> EntityWorldMut<'w> {
         Some(result)
     }
 
+    /// 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).
+    ///
+    /// # Panics
+    ///
+    /// If the entity has been despawned while this `EntityWorldMut` is still alive.
+    #[inline]
+    pub fn modify_component_by_id<R>(
+        &mut self,
+        component_id: ComponentId,
+        f: impl for<'a> FnOnce(MutUntyped<'a>) -> R,
+    ) -> Option<R> {
+        self.assert_not_despawned();
+
+        let result = self
+            .world
+            .modify_component_by_id(self.entity, component_id, 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`.
     ///
@@ -1326,6 +1455,23 @@ impl<'w> EntityWorldMut<'w> {
         unsafe { self.into_unsafe_entity_cell().get_mut() }
     }
 
+    /// Consumes `self` and gets mutable access to the component of type `T`
+    /// with the world `'w` lifetime for the current entity.
+    /// Returns `None` if the entity does not have a component of type `T`.
+    ///
+    /// # Panics
+    ///
+    /// If the entity has been despawned while this `EntityWorldMut` is still alive.
+    ///
+    /// # Safety
+    ///
+    /// - `T` must be a mutable component
+    #[inline]
+    pub unsafe fn into_mut_assume_mutable<T: Component>(self) -> Option<Mut<'w, T>> {
+        // SAFETY: consuming `self` implies exclusive access
+        unsafe { self.into_unsafe_entity_cell().get_mut_assume_mutable() }
+    }
+
     /// Gets a reference to the resource of the given type
     ///
     /// # Panics
@@ -1487,6 +1633,41 @@ impl<'w> EntityWorldMut<'w> {
         self.as_mutable().into_mut_by_id(component_ids)
     }
 
+    /// Returns [untyped mutable reference(s)](MutUntyped) to component(s) for
+    /// the current entity, based on the given [`ComponentId`]s.
+    /// Assumes the given [`ComponentId`]s refer to mutable components.
+    ///
+    /// **You should prefer to use the typed API [`EntityWorldMut::get_mut_assume_mutable`] where
+    /// possible and only use this in cases where the actual component types
+    /// are not known at compile time.**
+    ///
+    /// Unlike [`EntityWorldMut::get_mut_assume_mutable`], this returns untyped reference(s) to
+    /// component(s), and it's the job of the caller to ensure the correct
+    /// type(s) are dereferenced (if necessary).
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if the entity does
+    ///   not have a component.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component
+    ///   is requested multiple times.
+    ///
+    /// # Panics
+    ///
+    /// If the entity has been despawned while this `EntityWorldMut` is still alive.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the provided [`ComponentId`]s must refer to mutable components.
+    #[inline]
+    pub unsafe fn get_mut_assume_mutable_by_id<F: DynamicComponentFetch>(
+        &mut self,
+        component_ids: F,
+    ) -> Result<F::Mut<'_>, EntityComponentError> {
+        self.as_mutable()
+            .into_mut_assume_mutable_by_id(component_ids)
+    }
+
     /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped)
     /// to component(s) with lifetime `'w` for the current entity, based on the
     /// given [`ComponentId`]s.
@@ -1521,6 +1702,42 @@ impl<'w> EntityWorldMut<'w> {
         self.into_mutable().into_mut_by_id(component_ids)
     }
 
+    /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped)
+    /// to component(s) with lifetime `'w` for the current entity, based on the
+    /// given [`ComponentId`]s.
+    /// Assumes the given [`ComponentId`]s refer to mutable components.
+    ///
+    /// **You should prefer to use the typed API [`EntityWorldMut::into_mut_assume_mutable`] where
+    /// possible and only use this in cases where the actual component types
+    /// are not known at compile time.**
+    ///
+    /// Unlike [`EntityWorldMut::into_mut_assume_mutable`], this returns untyped reference(s) to
+    /// component(s), and it's the job of the caller to ensure the correct
+    /// type(s) are dereferenced (if necessary).
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if the entity does
+    ///   not have a component.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component
+    ///   is requested multiple times.
+    ///
+    /// # Panics
+    ///
+    /// If the entity has been despawned while this `EntityWorldMut` is still alive.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the provided [`ComponentId`]s must refer to mutable components.
+    #[inline]
+    pub unsafe fn into_mut_assume_mutable_by_id<F: DynamicComponentFetch>(
+        self,
+        component_ids: F,
+    ) -> Result<F::Mut<'w>, EntityComponentError> {
+        self.into_mutable()
+            .into_mut_assume_mutable_by_id(component_ids)
+    }
+
     /// Adds a [`Bundle`] of components to the entity.
     ///
     /// This will overwrite any previous value(s) of the same component type.
@@ -4396,6 +4613,26 @@ pub unsafe trait DynamicComponentFetch {
         self,
         cell: UnsafeEntityCell<'_>,
     ) -> Result<Self::Mut<'_>, EntityComponentError>;
+
+    /// Returns untyped mutable reference(s) to the component(s) with the
+    /// given [`ComponentId`]s, as determined by `self`.
+    /// Assumes all [`ComponentId`]s refer to mutable components.
+    ///
+    /// # Safety
+    ///
+    /// It is the caller's responsibility to ensure that:
+    /// - The given [`UnsafeEntityCell`] has mutable access to the fetched components.
+    /// - No other references to the fetched components exist at the same time.
+    /// - The requested components are all mutable.
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if a component is missing from the entity.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component is requested multiple times.
+    unsafe fn fetch_mut_assume_mutable(
+        self,
+        cell: UnsafeEntityCell<'_>,
+    ) -> Result<Self::Mut<'_>, EntityComponentError>;
 }
 
 // SAFETY:
@@ -4421,6 +4658,15 @@ unsafe impl DynamicComponentFetch for ComponentId {
         unsafe { cell.get_mut_by_id(self) }
             .map_err(|_| EntityComponentError::MissingComponent(self))
     }
+
+    unsafe fn fetch_mut_assume_mutable(
+        self,
+        cell: UnsafeEntityCell<'_>,
+    ) -> Result<Self::Mut<'_>, EntityComponentError> {
+        // SAFETY: caller ensures that the cell has mutable access to the component.
+        unsafe { cell.get_mut_assume_mutable_by_id(self) }
+            .map_err(|_| EntityComponentError::MissingComponent(self))
+    }
 }
 
 // SAFETY:
@@ -4443,6 +4689,13 @@ unsafe impl<const N: usize> DynamicComponentFetch for [ComponentId; N] {
     ) -> Result<Self::Mut<'_>, EntityComponentError> {
         <&Self>::fetch_mut(&self, cell)
     }
+
+    unsafe fn fetch_mut_assume_mutable(
+        self,
+        cell: UnsafeEntityCell<'_>,
+    ) -> Result<Self::Mut<'_>, EntityComponentError> {
+        <&Self>::fetch_mut_assume_mutable(&self, cell)
+    }
 }
 
 // SAFETY:
@@ -4497,6 +4750,34 @@ unsafe impl<const N: usize> DynamicComponentFetch for &'_ [ComponentId; N] {
 
         Ok(ptrs)
     }
+
+    unsafe fn fetch_mut_assume_mutable(
+        self,
+        cell: UnsafeEntityCell<'_>,
+    ) -> Result<Self::Mut<'_>, EntityComponentError> {
+        // Check for duplicate component IDs.
+        for i in 0..self.len() {
+            for j in 0..i {
+                if self[i] == self[j] {
+                    return Err(EntityComponentError::AliasedMutability(self[i]));
+                }
+            }
+        }
+
+        let mut ptrs = [const { MaybeUninit::uninit() }; N];
+        for (ptr, &id) in core::iter::zip(&mut ptrs, self) {
+            *ptr = MaybeUninit::new(
+                // SAFETY: caller ensures that the cell has mutable access to the component.
+                unsafe { cell.get_mut_assume_mutable_by_id(id) }
+                    .map_err(|_| EntityComponentError::MissingComponent(id))?,
+            );
+        }
+
+        // SAFETY: Each ptr was initialized in the loop above.
+        let ptrs = ptrs.map(|ptr| unsafe { MaybeUninit::assume_init(ptr) });
+
+        Ok(ptrs)
+    }
 }
 
 // SAFETY:
@@ -4543,6 +4824,30 @@ unsafe impl DynamicComponentFetch for &'_ [ComponentId] {
         }
         Ok(ptrs)
     }
+
+    unsafe fn fetch_mut_assume_mutable(
+        self,
+        cell: UnsafeEntityCell<'_>,
+    ) -> Result<Self::Mut<'_>, EntityComponentError> {
+        // Check for duplicate component IDs.
+        for i in 0..self.len() {
+            for j in 0..i {
+                if self[i] == self[j] {
+                    return Err(EntityComponentError::AliasedMutability(self[i]));
+                }
+            }
+        }
+
+        let mut ptrs = Vec::with_capacity(self.len());
+        for &id in self {
+            ptrs.push(
+                // SAFETY: caller ensures that the cell has mutable access to the component.
+                unsafe { cell.get_mut_assume_mutable_by_id(id) }
+                    .map_err(|_| EntityComponentError::MissingComponent(id))?,
+            );
+        }
+        Ok(ptrs)
+    }
 }
 
 // SAFETY:
@@ -4582,6 +4887,22 @@ unsafe impl DynamicComponentFetch for &'_ HashSet<ComponentId> {
         }
         Ok(ptrs)
     }
+
+    unsafe fn fetch_mut_assume_mutable(
+        self,
+        cell: UnsafeEntityCell<'_>,
+    ) -> Result<Self::Mut<'_>, EntityComponentError> {
+        let mut ptrs = HashMap::with_capacity_and_hasher(self.len(), Default::default());
+        for &id in self {
+            ptrs.insert(
+                id,
+                // SAFETY: caller ensures that the cell has mutable access to the component.
+                unsafe { cell.get_mut_assume_mutable_by_id(id) }
+                    .map_err(|_| EntityComponentError::MissingComponent(id))?,
+            );
+        }
+        Ok(ptrs)
+    }
 }
 
 #[cfg(test)]

crates/bevy_ecs/src/world/mod.rs

@@ -1332,6 +1332,35 @@ impl World {
         Ok(result)
     }
 
+    /// Temporarily removes a [`Component`] identified by the provided
+    /// [`ComponentId`] from the provided [`Entity`] and runs the provided
+    /// closure on it, returning the result if the component 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_by_id`](World::get_mut_by_id).
+    ///
+    /// You should prefer the typed [`modify_component`](World::modify_component)
+    /// whenever possible.
+    #[inline]
+    pub fn modify_component_by_id<R>(
+        &mut self,
+        entity: Entity,
+        component_id: ComponentId,
+        f: impl for<'a> FnOnce(MutUntyped<'a>) -> R,
+    ) -> Result<Option<R>, EntityMutableFetchError> {
+        let mut world = DeferredWorld::from(&mut *self);
+
+        let result = world.modify_component_by_id(entity, component_id, f)?;
+
+        self.flush();
+        Ok(result)
+    }
+
     /// Despawns the given [`Entity`], if it exists. This will also remove all of the entity's
     /// [`Components`](Component).
     ///

crates/bevy_ecs/src/world/unsafe_world_cell.rs

@@ -1081,6 +1081,54 @@ impl<'w> UnsafeEntityCell<'w> {
         }
     }
 
+    /// Retrieves a mutable untyped reference to the given `entity`'s [`Component`] of the given [`ComponentId`].
+    /// Returns `None` if the `entity` does not have a [`Component`] of the given type.
+    /// This method assumes the [`Component`] is mutable, skipping that check.
+    ///
+    /// **You should prefer to use the typed API [`UnsafeEntityCell::get_mut_assume_mutable`] where possible and only
+    /// use this in cases where the actual types are not known at compile time.**
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeEntityCell`] has permission to access the component mutably
+    /// - no other references to the component exist at the same time
+    /// - the component `T` is mutable
+    #[inline]
+    pub unsafe fn get_mut_assume_mutable_by_id(
+        self,
+        component_id: ComponentId,
+    ) -> Result<MutUntyped<'w>, GetEntityMutByIdError> {
+        self.world.assert_allows_mutable_access();
+
+        let info = self
+            .world
+            .components()
+            .get_info(component_id)
+            .ok_or(GetEntityMutByIdError::InfoNotFound)?;
+
+        // SAFETY: entity_location is valid, component_id is valid as checked by the line above
+        unsafe {
+            get_component_and_ticks(
+                self.world,
+                component_id,
+                info.storage_type(),
+                self.entity,
+                self.location,
+            )
+            .map(|(value, cells, caller)| MutUntyped {
+                // SAFETY: world access validated by caller and ties world lifetime to `MutUntyped` lifetime
+                value: value.assert_unique(),
+                ticks: TicksMut::from_tick_cells(
+                    cells,
+                    self.world.last_change_tick(),
+                    self.world.change_tick(),
+                ),
+                changed_by: caller.map(|caller| caller.deref_mut()),
+            })
+            .ok_or(GetEntityMutByIdError::ComponentNotFound)
+        }
+    }
+
     /// Returns the source code location from which this entity has been spawned.
     pub fn spawned_by(self) -> MaybeLocation {
         self.world()