diff --git a/crates/bevy_ecs/macros/src/world_query.rs b/crates/bevy_ecs/macros/src/world_query.rs index 5a7d164b80..1c92ef9705 100644 --- a/crates/bevy_ecs/macros/src/world_query.rs +++ b/crates/bevy_ecs/macros/src/world_query.rs @@ -168,6 +168,17 @@ pub(crate) fn world_query_impl( }) } + #[inline] + fn apply(state: &mut Self::State, system_meta: &#path::system::SystemMeta, world: &mut #path::world::World) { + #(<#field_types as #path::query::WorldQuery>::apply(&mut state.#named_field_idents, system_meta, world);)* + } + + + #[inline] + fn queue(state: &mut Self::State, system_meta: &#path::system::SystemMeta, mut world: #path::world::DeferredWorld) { + #(<#field_types as #path::query::WorldQuery>::queue(&mut state.#named_field_idents, system_meta, world.reborrow());)* + } + fn matches_component_set(state: &Self::State, _set_contains_id: &impl Fn(#path::component::ComponentId) -> bool) -> bool { true #(&& <#field_types>::matches_component_set(&state.#named_field_idents, _set_contains_id))* } diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 2564223972..215a01cac2 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -6,9 +6,10 @@ use crate::{ entity::{Entities, Entity, EntityLocation}, query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery}, storage::{ComponentSparseSet, Table, TableRow}, + system::SystemMeta, world::{ - unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept, - FilteredEntityMut, FilteredEntityRef, Mut, Ref, World, + unsafe_world_cell::UnsafeWorldCell, DeferredWorld, EntityMut, EntityMutExcept, EntityRef, + EntityRefExcept, FilteredEntityMut, FilteredEntityRef, Mut, Ref, World, }, }; use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref}; @@ -2278,6 +2279,16 @@ unsafe impl WorldQuery for Option { ) -> bool { true } + + const HAS_DEFERRED: bool = T::HAS_DEFERRED; + + fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) { + ::apply(state, system_meta, world); + } + + fn queue(state: &mut Self::State, system_meta: &SystemMeta, world: DeferredWorld) { + ::queue(state, system_meta, world); + } } // SAFETY: defers to soundness of `T: WorldQuery` impl @@ -2489,6 +2500,268 @@ impl ReleaseStateQueryData for Has { } } +// gate `DeferredMut` behind support for Parallel +bevy_utils::cfg::parallel! { + use core::ops::{Deref, DerefMut}; + use crate::entity::EntityHashMap; + use bevy_utils::Parallel; + + /// Provides "fake" mutable access to the component `T` + /// + /// `DeferredMut` only accesses `&T` from the world, but when mutably + /// dereferenced will clone it and return a reference to that cloned value. + /// Once the `DeferredMut` is dropped, the query keeps track of the new value + /// and inserts it into the world at the next sync point. + /// + /// This can be used to "mutably" access immutable components! + /// However, this will still be slower than direct mutation, so this should + /// mainly be used for its ergonomics. + /// + /// # Examples + /// + /// ``` + /// # use bevy_ecs::component::Component; + /// # use bevy_ecs::query::DeferredMut; + /// # use bevy_ecs::query::With; + /// # use bevy_ecs::system::IntoSystem; + /// # use bevy_ecs::system::Query; + /// # + /// # #[derive(Component)] + /// # struct Poisoned; + /// # + /// #[derive(Component, Clone)] + /// #[component(immutable)] + /// struct Health(u32); + /// + /// fn tick_poison(mut health_query: Query, With>) { + /// for mut health in &health_query { + /// health.0 -= 1; + /// } + /// } + /// # bevy_ecs::system::assert_is_system(tick_poison); + /// ``` + /// + /// # Footguns + /// + /// 1. The mutations tracked by `DeferredMut` will *not* be applied if used + /// through the manual [`QueryState`](super::QueryState) API. Instead, it + /// should be used through a query in a system param or [`SystemState`](crate::system::SystemState). + /// + /// 2. It's possible to query multiple `DeferredMut` values from the same entity. + /// However, since mutations are deferred, each new value won't see the changes + /// applied to previous iterations. + /// + /// Normally, the final iteration will be the one that "wins" and gets inserted + /// onto the entity, but parallelism can mess with that, too. Since `DeferredMut` + /// internally uses a thread-local [`EntityHashMap`] to keep track of mutations, + /// if two `DeferredMut` values for the same entity are created in the same system + /// on different threads, then they'll each be inserted into the entity in an + /// undetermined order. + pub struct DeferredMut<'w, 's, T: Component> { + entity: Entity, + old: &'w T, + new: Option, + record: &'s DeferredMutations, + } + + impl<'w, 's, T: Component> DeferredMut<'w, 's, T> { + /// Returns a reference to the `T` value still present in the ECS + #[inline] + pub fn stale(&self) -> &'w T { + self.old + } + + /// Returns a reference to the `T` value currently being updated. + /// If none is present yet, this method will clone from `Self::stale` + #[inline] + pub fn fresh(&mut self) -> &mut T where T: Clone { + self.get_fresh_or_insert(self.old.clone()) + } + + /// Returns a (possibly absent) reference to the `T` value currently being updated. + #[inline] + pub fn get_fresh(&mut self) -> Option<&mut T> { + self.new.as_mut() + } + + /// Returns a reference to the `T` value currently being updated. + /// If absent, it will insert the provided value. + #[inline] + pub fn get_fresh_or_insert(&mut self, value: T) -> &mut T { + self.new.get_or_insert(value) + } + + /// Replaces the `T` value currently being updated + #[inline] + pub fn insert(&mut self, value: T) { + self.new = Some(value); + } + } + + impl<'w, 's, T: Component> Drop for DeferredMut<'w, 's, T> { + #[inline] + fn drop(&mut self) { + if let Some(new) = self.new.take() { + self.record.insert(self.entity, new); + } + } + } + + impl<'w, 's, T: Component> Deref for DeferredMut<'w, 's, T> { + type Target = T; + + #[inline] + fn deref(&self) -> &Self::Target { + self.new.as_ref().unwrap_or(self.old) + } + } + + impl<'w, 's, T: Component + Clone> DerefMut for DeferredMut<'w, 's, T> { + #[inline] + fn deref_mut(&mut self) -> &mut Self::Target { + self.fresh() + } + } + + /// The [`WorldQuery::State`] type for [`DeferredMut`] + pub struct DeferredMutState { + internal: <&'static T as WorldQuery>::State, + record: DeferredMutations, + } + + struct DeferredMutations(Parallel>); + + impl Default for DeferredMutations { + fn default() -> Self { + Self(Default::default()) + } + } + + impl DeferredMutations { + #[inline] + fn insert(&self, entity: Entity, component: T) { + self.0.scope(|map| map.insert(entity, component)); + } + + #[inline] + fn drain(&mut self) -> impl Iterator { + self.0.drain() + } + } + + // SAFETY: impl defers to `<&T as WorldQuery>` for all methods + unsafe impl<'__w, '__s, T: Component> WorldQuery for DeferredMut<'__w, '__s, T> { + type Fetch<'w> = ReadFetch<'w, T>; + + type State = DeferredMutState; + + fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { + fetch + } + + unsafe fn init_fetch<'w, 's>( + world: UnsafeWorldCell<'w>, + state: &'s Self::State, + last_run: Tick, + this_run: Tick, + ) -> Self::Fetch<'w> { + // SAFETY: invariants are upheld by the caller + unsafe { <&T as WorldQuery>::init_fetch(world, &state.internal, last_run, this_run) } + } + + const IS_DENSE: bool = <&T as WorldQuery>::IS_DENSE; + + unsafe fn set_archetype<'w>( + fetch: &mut Self::Fetch<'w>, + state: &Self::State, + archetype: &'w Archetype, + table: &'w Table, + ) { + <&T as WorldQuery>::set_archetype(fetch, &state.internal, archetype, table); + } + + unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) { + <&T as WorldQuery>::set_table(fetch, &state.internal, table); + } + + fn update_component_access(state: &Self::State, access: &mut FilteredAccess) { + <&T as WorldQuery>::update_component_access(&state.internal, access); + } + + fn init_state(world: &mut World) -> Self::State { + DeferredMutState { + internal: <&T as WorldQuery>::init_state(world), + record: Default::default(), + } + } + + fn get_state(components: &Components) -> Option { + Some(DeferredMutState { + internal: <&T as WorldQuery>::get_state(components)?, + record: Default::default(), + }) + } + + fn matches_component_set( + state: &Self::State, + set_contains_id: &impl Fn(ComponentId) -> bool, + ) -> bool { + <&T as WorldQuery>::matches_component_set(&state.internal, set_contains_id) + } + + const HAS_DEFERRED: bool = true; + + fn apply(state: &mut Self::State, _system_meta: &SystemMeta, world: &mut World) { + world.insert_batch(state.record.drain()); + } + + fn queue(state: &mut Self::State, _system_meta: &SystemMeta, mut world: DeferredWorld) { + world + .commands() + .insert_batch(state.record.drain().collect::>()); + } + } + + // SAFETY: DeferredMut defers to &T internally, so it must be readonly and Self::ReadOnly = Self. + unsafe impl<'__w, '__s, T: Component> QueryData for DeferredMut<'__w, '__s, T> { + const IS_READ_ONLY: bool = true; + + type ReadOnly = Self; + + type Item<'w, 's> = DeferredMut<'w, 's, T>; + + fn shrink<'wlong: 'wshort, 'wshort, 's>( + item: Self::Item<'wlong, 's>, + ) -> Self::Item<'wshort, 's> { + item + } + + unsafe fn fetch<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entity: Entity, + table_row: TableRow, + ) -> Self::Item<'w, 's> { + // SAFETY: invariants are upheld by the caller + let old = + unsafe { <&T as QueryData>::fetch(&state.internal, fetch, entity, table_row) }; + DeferredMut { + entity, + old, + // NOTE: we could try to get an existing updated component from the record, + // but we can't reliably do that across all threads. Better to say that all + // newly-created DeferredMut values will match what's in the ECS. + new: None, + record: &state.record, + } + } + } + + // SAFETY: Tracked only accesses &T from the world. Though it provides mutable access, it only + // applies those changes through commands. + unsafe impl<'__w, '__s, T: Component> ReadOnlyQueryData for DeferredMut<'__w, '__s, T> {} +} + /// The `AnyOf` query parameter fetches entities with any of the component types included in T. /// /// `Query>` is equivalent to `Query<(Option<&A>, Option<&B>, Option<&mut C>), Or<(With, With, With)>>`. @@ -2586,6 +2859,10 @@ macro_rules! impl_anytuple_fetch { unused_variables, reason = "Zero-length tuples won't use any of the parameters." )] + #[allow( + unused_mut, + reason = "Zero-length tuples won't access any of the parameters mutably." + )] #[allow( clippy::unused_unit, reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." @@ -2683,6 +2960,16 @@ macro_rules! impl_anytuple_fetch { let ($($name,)*) = _state; false $(|| $name::matches_component_set($name, _set_contains_id))* } + + const HAS_DEFERRED: bool = false $(|| $name::HAS_DEFERRED)*; + + fn apply(($($state,)*): &mut Self::State, system_meta: &SystemMeta, world: &mut World) { + $(<$name as WorldQuery>::apply($state, system_meta, world);)* + } + + fn queue(($($state,)*): &mut Self::State, system_meta: &SystemMeta, mut world: DeferredWorld) { + $(<$name as WorldQuery>::queue($state, system_meta, world.reborrow());)* + } } #[expect( diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index f9f4861b79..c0d32294d2 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -4,7 +4,8 @@ use crate::{ entity::{Entities, Entity}, query::{DebugCheckedUnwrap, FilteredAccess, StorageSwitch, WorldQuery}, storage::{ComponentSparseSet, Table, TableRow}, - world::{unsafe_world_cell::UnsafeWorldCell, World}, + system::SystemMeta, + world::{unsafe_world_cell::UnsafeWorldCell, DeferredWorld, World}, }; use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref}; use bevy_utils::prelude::DebugName; @@ -378,6 +379,10 @@ macro_rules! impl_or_query_filter { unused_variables, reason = "Zero-length tuples won't use any of the parameters." )] + #[allow( + unused_mut, + reason = "Zero-length tuples won't access any of the parameters mutably." + )] #[allow( clippy::unused_unit, reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." @@ -478,6 +483,16 @@ macro_rules! impl_or_query_filter { let ($($filter,)*) = state; false $(|| $filter::matches_component_set($filter, set_contains_id))* } + + const HAS_DEFERRED: bool = false $(|| $filter::HAS_DEFERRED)*; + + fn apply(($($state,)*): &mut Self::State, system_meta: &SystemMeta, world: &mut World) { + $(<$filter as WorldQuery>::apply($state, system_meta, world);)* + } + + fn queue(($($state,)*): &mut Self::State, system_meta: &SystemMeta, mut world: DeferredWorld) { + $(<$filter as WorldQuery>::queue($state, system_meta, world.reborrow());)* + } } #[expect( diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index 1c739927ac..669c271d24 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/crates/bevy_ecs/src/query/world_query.rs @@ -3,7 +3,8 @@ use crate::{ component::{ComponentId, Components, Tick}, query::FilteredAccess, storage::Table, - world::{unsafe_world_cell::UnsafeWorldCell, World}, + system::SystemMeta, + world::{unsafe_world_cell::UnsafeWorldCell, DeferredWorld, World}, }; use variadics_please::all_tuples; @@ -121,6 +122,34 @@ pub unsafe trait WorldQuery { /// access to [`Components`]. fn get_state(components: &Components) -> Option; + /// Returns true if (and only if) this [`WorldQuery`] contains deferred state + /// that needs to be applied at the next sync point. If this is set to `false`, + /// `apply` or `queue` may not be called. + const HAS_DEFERRED: bool = false; + + /// Applies any deferred mutations stored in this [`WorldQuery`]'s state. + /// This is used to apply [`Commands`] during [`ApplyDeferred`](crate::prelude::ApplyDeferred). + /// + /// If this is not a no-op, then `HAS_DEFERRED` should be `true` + /// + /// [`Commands`]: crate::prelude::Commands + #[inline] + #[expect( + unused_variables, + reason = "The parameters here are intentionally unused by the default implementation; however, putting underscores here will result in the underscores being copied by rust-analyzer's tab completion." + )] + fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) {} + + /// Queues any deferred mutations to be applied at the next [`ApplyDeferred`](crate::prelude::ApplyDeferred). + /// + /// If this is not a no-op, then `HAS_DEFERRED` should be `true` + #[inline] + #[expect( + unused_variables, + reason = "The parameters here are intentionally unused by the default implementation; however, putting underscores here will result in the underscores being copied by rust-analyzer's tab completion." + )] + fn queue(state: &mut Self::State, system_meta: &SystemMeta, world: DeferredWorld) {} + /// Returns `true` if this query matches a set of components. Otherwise, returns `false`. /// /// Used to check which [`Archetype`]s can be skipped by the query @@ -211,6 +240,22 @@ macro_rules! impl_tuple_world_query { Some(($($name::get_state(components)?,)*)) } + const HAS_DEFERRED: bool = false $(|| $name::HAS_DEFERRED)*; + + #[inline] + fn apply(($($name,)*): &mut Self::State, system_meta: &SystemMeta, world: &mut World) { + $($name::apply($name, system_meta, world);)* + } + + #[inline] + #[allow( + unused_mut, + reason = "The `world` parameter is unused for zero-length tuples; however, it must be mutable for other lengths of tuples." + )] + fn queue(($($name,)*): &mut Self::State, system_meta: &SystemMeta, mut world: DeferredWorld) { + $($name::queue($name, system_meta, world.reborrow());)* + } + fn matches_component_set(state: &Self::State, set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { let ($($name,)*) = state; true $(&& $name::matches_component_set($name, set_contains_id))* diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs index d9b433e7f2..df46aee648 100644 --- a/crates/bevy_ecs/src/system/system_param.rs +++ b/crates/bevy_ecs/src/system/system_param.rs @@ -355,6 +355,10 @@ unsafe impl SystemParam for Qu world, ); component_access_set.add(state.component_access.clone()); + + if D::HAS_DEFERRED || F::HAS_DEFERRED { + system_meta.set_has_deferred(); + } } #[inline] @@ -370,6 +374,18 @@ unsafe impl SystemParam for Qu // The caller ensures the world matches the one used in init_state. unsafe { state.query_unchecked_with_ticks(world, system_meta.last_run, change_tick) } } + + #[inline] + fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) { + D::apply(&mut state.fetch_state, system_meta, world); + F::apply(&mut state.filter_state, system_meta, world); + } + + #[inline] + fn queue(state: &mut Self::State, system_meta: &SystemMeta, mut world: DeferredWorld) { + D::queue(&mut state.fetch_state, system_meta, world.reborrow()); + F::queue(&mut state.filter_state, system_meta, world); + } } fn assert_component_access_compatibility( diff --git a/release-content/release-notes/new_querydata.md b/release-content/release-notes/new_querydata.md new file mode 100644 index 0000000000..ad94e56439 --- /dev/null +++ b/release-content/release-notes/new_querydata.md @@ -0,0 +1,45 @@ +--- +title: New QueryData Types +authors: ["@ecoskey"] +pull_requests: [19602] +--- + +Bevy queries have some new powers for advanced users. Namely, custom `WorldQuery` +implementations can store and apply "deferred" mutations, just like `Commands`! +This release includes a few new types making use of this capability, and +we're sure third-party crates will find all kinds of new ways to do cool stuff +with this. + +## `DeferredMut` + +When working with immutable components in Bevy, the acts of reading and writing +component values are very clearly separated. This can be valuable, especially +if a component has expensive hooks or observers attached and `insert`ing it +has a significant cost, but in some cases it can feel like boilerplate. + +`DeferredMut` is meant to improve the ergonomics of the latter case, by providing +"fake" mutable access to any component, even immutable ones! Internally, it +keeps track of any modifications and inserts them into the ECS at the next +sync point. + +```rs +// without `DeferredMut` +pub fn tick_poison( + mut commands: Commands, + query: Query<(Entity, &Health), With> +) { + for (entity, Health(health_points)) in query { + commands.insert(entity, Health(health_points - 1)) + } +} + +// with `DeferredMut` +pub fn tick_poison( + mut health_query: Query, With> +) { + for mut health in &health_query { + health.0 -= 1; + } +} + +```