From 01bde1ff36ed6f6ae07de2c0e5452733a5fa043c Mon Sep 17 00:00:00 2001 From: Zac Harrold Date: Fri, 6 Dec 2024 09:24:05 +1100 Subject: [PATCH 1/6] Added `with_component` to `EntityWorldMut`, `DeferredWorld`, and `World` --- crates/bevy_ecs/src/world/deferred_world.rs | 94 ++++++++++++++- crates/bevy_ecs/src/world/entity_ref.rs | 123 ++++++++++++++++++++ crates/bevy_ecs/src/world/mod.rs | 38 ++++++ 3 files changed, 254 insertions(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/world/deferred_world.rs b/crates/bevy_ecs/src/world/deferred_world.rs index 54bacb13fd47d..996471f434356 100644 --- a/crates/bevy_ecs/src/world/deferred_world.rs +++ b/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,98 @@ 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. + /// + /// While this is available for all components, it's recommended to only be + /// used with immutable components. + /// When available, prefer using [`get_mut`](DeferredWorld::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::(); + /// # + /// # let entity = world.spawn(Foo(false)).id(); + /// # + /// # let mut world = DeferredWorld::from(&mut world); + /// # + /// world.with_component(entity, |foo: &mut Foo| { + /// foo.0 = true; + /// }); + /// # + /// # assert_eq!(world.get::(entity), Some(&Foo(true))); + /// ``` + #[inline] + pub fn with_component( + &mut self, + entity: Entity, + f: impl FnOnce(&mut T) -> R, + ) -> Result, 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::() else { + return Ok(None); + }; + + let entity_cell = self.get_entity_mut(entity)?; + + if !entity_cell.contains::() { + 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::() + .expect("component access confirmed above") + }; + + let result = f(&mut component); + + // 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 diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 9ee25243bf624..f8a4dd37337bc 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -1099,6 +1099,46 @@ 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. + /// + /// While this is available for all components, it's recommended to only be + /// used with immutable components. + /// When available, prefer using [`get_mut`](EntityMut::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::(); + /// # + /// # let entity = world.spawn(Foo(false)).id(); + /// # + /// # let mut entity = world.entity_mut(entity); + /// # + /// # assert_eq!(entity.get::(), Some(&Foo(false))); + /// # + /// entity.with_component(|foo: &mut Foo| { + /// foo.0 = true; + /// }); + /// # + /// # assert_eq!(entity.get::(), Some(&Foo(true))); + /// ``` + #[inline] + pub fn with_component(&mut self, f: impl FnOnce(&mut T) -> R) -> Option { + self.world + .with_component(self.entity, f) + .expect("entity access must be valid") + } + /// 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`. /// @@ -4922,4 +4962,87 @@ mod tests { world.flush(); assert_eq!(world.resource_mut::().0.as_slice(), &expected[..]); } + + #[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::(); + world + .register_component_hooks::() + .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::(), Some(&Foo(false))); + + entity.with_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::(), 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); + } } diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 491c3648dd043..e0c84f3624726 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -1482,6 +1482,44 @@ 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. + /// + /// While this is available for all components, it's recommended to only be + /// used with immutable components. + /// When available, 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::(); + /// # + /// # let entity = world.spawn(Foo(false)).id(); + /// # + /// world.with_component(entity, |foo: &mut Foo| { + /// foo.0 = true; + /// }); + /// # + /// # assert_eq!(world.get::(entity), Some(&Foo(true))); + /// ``` + #[inline] + pub fn with_component( + &mut self, + entity: Entity, + f: impl FnOnce(&mut T) -> R, + ) -> Result, EntityFetchError> { + DeferredWorld::from(self).with_component(entity, f) + } + /// 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. From 4241af48ccc9d0b2c70bbf7f2fe60f41d79a279f Mon Sep 17 00:00:00 2001 From: Zac Harrold Date: Fri, 6 Dec 2024 10:22:19 +1100 Subject: [PATCH 2/6] Fix docs --- crates/bevy_ecs/src/world/deferred_world.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/world/deferred_world.rs b/crates/bevy_ecs/src/world/deferred_world.rs index 996471f434356..823911b8a1edd 100644 --- a/crates/bevy_ecs/src/world/deferred_world.rs +++ b/crates/bevy_ecs/src/world/deferred_world.rs @@ -94,7 +94,7 @@ impl<'w> DeferredWorld<'w> { /// # Examples /// /// ```rust - /// # use bevy_ecs::prelude::*; + /// # use bevy_ecs::{prelude::*, world::DeferredWorld}; /// # /// #[derive(Component, PartialEq, Eq, Debug)] /// #[component(immutable)] From 0c22ba8d4e4d47cdd9bedc3ba8917d5f18b9fbb2 Mon Sep 17 00:00:00 2001 From: Zac Harrold Date: Fri, 13 Dec 2024 13:07:00 +1100 Subject: [PATCH 3/6] Ensure not despawned and flush afterwards if required. --- crates/bevy_ecs/src/world/entity_ref.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index bc4046fb89d28..a7138a577ff9c 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -1132,11 +1132,23 @@ impl<'w> EntityWorldMut<'w> { /// # /// # assert_eq!(entity.get::(), Some(&Foo(true))); /// ``` + /// + /// # Panics + /// + /// If the entity has been despawned while this `EntityWorldMut` is still alive. #[inline] pub fn with_component(&mut self, f: impl FnOnce(&mut T) -> R) -> Option { - self.world + self.assert_not_despawned(); + + let result = self + .world .with_component(self.entity, f) - .expect("entity access must be valid") + .expect("entity access must be valid")?; + + self.world.flush(); + self.update_location(); + + Some(result) } /// Gets mutable access to the component of type `T` for the current entity. From 1cb379d833fdb664f2a234e628b80754d3cbb109 Mon Sep 17 00:00:00 2001 From: Zac Harrold Date: Fri, 13 Dec 2024 13:17:46 +1100 Subject: [PATCH 4/6] Simulate adding when using `with_component` --- crates/bevy_ecs/src/world/deferred_world.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/bevy_ecs/src/world/deferred_world.rs b/crates/bevy_ecs/src/world/deferred_world.rs index ab8ef0cd4346d..8f89b340a7797 100644 --- a/crates/bevy_ecs/src/world/deferred_world.rs +++ b/crates/bevy_ecs/src/world/deferred_world.rs @@ -158,6 +158,9 @@ impl<'w> DeferredWorld<'w> { 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. From 16dce7986b4d6c62b84cc31719d5c90a5b0e27e1 Mon Sep 17 00:00:00 2001 From: Zac Harrold Date: Fri, 13 Dec 2024 13:32:53 +1100 Subject: [PATCH 5/6] Make `DeferredWorld::with_component` `pub(crate)` Also ensure commands are always flushed when using `World::with_component` --- crates/bevy_ecs/src/world/deferred_world.rs | 2 +- crates/bevy_ecs/src/world/entity_ref.rs | 1 - crates/bevy_ecs/src/world/mod.rs | 6 +++++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/bevy_ecs/src/world/deferred_world.rs b/crates/bevy_ecs/src/world/deferred_world.rs index 8f89b340a7797..b7b669ae2da07 100644 --- a/crates/bevy_ecs/src/world/deferred_world.rs +++ b/crates/bevy_ecs/src/world/deferred_world.rs @@ -114,7 +114,7 @@ impl<'w> DeferredWorld<'w> { /// # assert_eq!(world.get::(entity), Some(&Foo(true))); /// ``` #[inline] - pub fn with_component( + pub(crate) fn with_component( &mut self, entity: Entity, f: impl FnOnce(&mut T) -> R, diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index a7138a577ff9c..e0ae904230098 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -1145,7 +1145,6 @@ impl<'w> EntityWorldMut<'w> { .with_component(self.entity, f) .expect("entity access must be valid")?; - self.world.flush(); self.update_location(); Some(result) diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 0789dc603b8bd..4677501f3be1b 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -1516,7 +1516,11 @@ impl World { entity: Entity, f: impl FnOnce(&mut T) -> R, ) -> Result, EntityFetchError> { - DeferredWorld::from(self).with_component(entity, f) + let result = DeferredWorld::from(&mut *self).with_component(entity, f)?; + + self.flush(); + + Ok(result) } /// Despawns the given `entity`, if it exists. This will also remove all of the entity's From 136eae4eafd83d1ab79d9b6972f02f67e600efd6 Mon Sep 17 00:00:00 2001 From: Zac Harrold Date: Fri, 13 Dec 2024 14:29:01 +1100 Subject: [PATCH 6/6] Removed outdated example --- crates/bevy_ecs/src/world/deferred_world.rs | 23 --------------------- 1 file changed, 23 deletions(-) diff --git a/crates/bevy_ecs/src/world/deferred_world.rs b/crates/bevy_ecs/src/world/deferred_world.rs index b7b669ae2da07..0954125099664 100644 --- a/crates/bevy_ecs/src/world/deferred_world.rs +++ b/crates/bevy_ecs/src/world/deferred_world.rs @@ -90,29 +90,6 @@ impl<'w> DeferredWorld<'w> { /// While this is available for all components, it's recommended to only be /// used with immutable components. /// When available, prefer using [`get_mut`](DeferredWorld::get_mut). - /// - /// # Examples - /// - /// ```rust - /// # use bevy_ecs::{prelude::*, world::DeferredWorld}; - /// # - /// #[derive(Component, PartialEq, Eq, Debug)] - /// #[component(immutable)] - /// struct Foo(bool); - /// - /// # let mut world = World::default(); - /// # world.register_component::(); - /// # - /// # let entity = world.spawn(Foo(false)).id(); - /// # - /// # let mut world = DeferredWorld::from(&mut world); - /// # - /// world.with_component(entity, |foo: &mut Foo| { - /// foo.0 = true; - /// }); - /// # - /// # assert_eq!(world.get::(entity), Some(&Foo(true))); - /// ``` #[inline] pub(crate) fn with_component( &mut self,