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>
This commit is contained in:
Zachary Harrold 2024-12-31 11:23:44 +11:00 committed by GitHub
parent 4460a4d9ed
commit 3f19997096
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 272 additions and 1 deletions

View File

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

View File

@ -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);
}
}

View File

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