From 921ff6701f17cdddd37f345255e779b4a277c867 Mon Sep 17 00:00:00 2001 From: Zachary Harrold Date: Wed, 26 Mar 2025 07:52:07 +1100 Subject: [PATCH] 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 | 35 +- crates/bevy_ecs/src/world/entity_ref.rs | 321 ++++++++++++++++++ crates/bevy_ecs/src/world/mod.rs | 29 ++ .../bevy_ecs/src/world/unsafe_world_cell.rs | 48 +++ 4 files changed, 430 insertions(+), 3 deletions(-) diff --git a/crates/bevy_ecs/src/world/deferred_world.rs b/crates/bevy_ecs/src/world/deferred_world.rs index 4332d12163..d80f838108 100644 --- a/crates/bevy_ecs/src/world/deferred_world.rs +++ b/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::() }; + + 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( + &mut self, + entity: Entity, + component_id: ComponentId, + f: impl for<'a> FnOnce(MutUntyped<'a>) -> R, + ) -> Result, EntityMutableFetchError> { let entity_cell = self.get_entity_mut(entity)?; - if !entity_cell.contains::() { + 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::() + .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; diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 3813d34745..7c1a14dcf0 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/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( + &mut self, + component_ids: F, + ) -> Result, 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( + &self, + component_ids: F, + ) -> Result, 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( + self, + component_ids: F, + ) -> Result, 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( + &mut self, + component_id: ComponentId, + f: impl for<'a> FnOnce(MutUntyped<'a>) -> R, + ) -> Option { + 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(self) -> Option> { + // 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( + &mut self, + component_ids: F, + ) -> Result, 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( + self, + component_ids: F, + ) -> Result, 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, 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, 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, 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 DynamicComponentFetch for [ComponentId; N] { ) -> Result, EntityComponentError> { <&Self>::fetch_mut(&self, cell) } + + unsafe fn fetch_mut_assume_mutable( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + <&Self>::fetch_mut_assume_mutable(&self, cell) + } } // SAFETY: @@ -4497,6 +4750,34 @@ unsafe impl DynamicComponentFetch for &'_ [ComponentId; N] { Ok(ptrs) } + + unsafe fn fetch_mut_assume_mutable( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, 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, 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 { } Ok(ptrs) } + + unsafe fn fetch_mut_assume_mutable( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, 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)] diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 5f2da150c7..6edc5186c7 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/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( + &mut self, + entity: Entity, + component_id: ComponentId, + f: impl for<'a> FnOnce(MutUntyped<'a>) -> R, + ) -> Result, 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). /// diff --git a/crates/bevy_ecs/src/world/unsafe_world_cell.rs b/crates/bevy_ecs/src/world/unsafe_world_cell.rs index f094906b12..8bca668ffc 100644 --- a/crates/bevy_ecs/src/world/unsafe_world_cell.rs +++ b/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, 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()