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.
This commit is contained in:
Zachary Harrold 2025-03-26 07:52:07 +11:00 committed by François Mockers
parent 8dc1a7cd81
commit 9705eaef02
4 changed files with 430 additions and 3 deletions

View File

@ -101,9 +101,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);
}
@ -140,11 +169,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;

View File

@ -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)]

View File

@ -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).
///

View File

@ -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()