From bea0a0a9bc37c66c8495bdee7247fc53dbbf07f1 Mon Sep 17 00:00:00 2001 From: Chris Russell <8494645+chescock@users.noreply.github.com> Date: Mon, 5 May 2025 19:23:46 -0400 Subject: [PATCH] Let `FilteredEntity(Ref|Mut)` receive access when nested. (#18236) # Objective Let `FilteredEntityRef` and `FilteredEntityMut` receive access when nested inside tuples or `#[derive(QueryData)]` types. Make sure to exclude any access that would conflict with other subqueries! Fixes #14349 ## Solution Replace `WorldQuery::set_access(state, access)` with a new method, `QueryData::provide_extra_access(state, access, available_access)`, that passes both the total available access and the currently used access. This is called after `WorldQuery::update_component_access()`, so any access used by ordinary subqueries will be known. `FilteredEntityRef` and `FilteredEntityMut` can use the combination to determine how much access they can safely take, while tuples can safely pass those parameters directly to their subqueries. This requires a new `Access::remove_conflicting_access()` method that can be used to remove any access that would conflict with existing access. Implementing this method was easier by first factoring some common set manipulation code out of `Access::extend`. I can extract that refactoring to a separate PR if desired. Have `FilteredEntity(Ref|Mut)` store `Access` instead of `FilteredAccess` because they do not need to keep track of the filter. This was necessary in an early draft but no longer is. I left it in because it's small and I'm touching that code anyway, but I can extract it to a separate PR if desired. --- crates/bevy_ecs/macros/src/query_data.rs | 16 ++ crates/bevy_ecs/src/query/access.rs | 247 +++++++++++++++++------ crates/bevy_ecs/src/query/builder.rs | 92 ++++++++- crates/bevy_ecs/src/query/fetch.rs | 101 ++++++--- crates/bevy_ecs/src/query/state.rs | 53 +++-- crates/bevy_ecs/src/query/world_query.rs | 14 +- crates/bevy_ecs/src/system/query.rs | 10 +- crates/bevy_ecs/src/world/entity_ref.rs | 16 -- 8 files changed, 407 insertions(+), 142 deletions(-) diff --git a/crates/bevy_ecs/macros/src/query_data.rs b/crates/bevy_ecs/macros/src/query_data.rs index d919d0b05e..4e4529e631 100644 --- a/crates/bevy_ecs/macros/src/query_data.rs +++ b/crates/bevy_ecs/macros/src/query_data.rs @@ -268,6 +268,14 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { } } + fn provide_extra_access( + state: &mut Self::State, + access: &mut #path::query::Access<#path::component::ComponentId>, + available_access: &#path::query::Access<#path::component::ComponentId>, + ) { + #(<#field_types>::provide_extra_access(&mut state.#named_field_idents, access, available_access);)* + } + /// SAFETY: we call `fetch` for each member that implements `Fetch`. #[inline(always)] unsafe fn fetch<'__w>( @@ -305,6 +313,14 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { } } + fn provide_extra_access( + state: &mut Self::State, + access: &mut #path::query::Access<#path::component::ComponentId>, + available_access: &#path::query::Access<#path::component::ComponentId>, + ) { + #(<#field_types>::provide_extra_access(&mut state.#named_field_idents, access, available_access);)* + } + /// SAFETY: we call `fetch` for each member that implements `Fetch`. #[inline(always)] unsafe fn fetch<'__w>( diff --git a/crates/bevy_ecs/src/query/access.rs b/crates/bevy_ecs/src/query/access.rs index d8d9c8c3fa..9c63cb5a74 100644 --- a/crates/bevy_ecs/src/query/access.rs +++ b/crates/bevy_ecs/src/query/access.rs @@ -431,78 +431,58 @@ impl Access { /// Adds all access from `other`. pub fn extend(&mut self, other: &Access) { - let component_read_and_writes_inverted = - self.component_read_and_writes_inverted || other.component_read_and_writes_inverted; - let component_writes_inverted = - self.component_writes_inverted || other.component_writes_inverted; - - match ( - self.component_read_and_writes_inverted, + invertible_union_with( + &mut self.component_read_and_writes, + &mut self.component_read_and_writes_inverted, + &other.component_read_and_writes, other.component_read_and_writes_inverted, - ) { - (true, true) => { - self.component_read_and_writes - .intersect_with(&other.component_read_and_writes); - } - (true, false) => { - self.component_read_and_writes - .difference_with(&other.component_read_and_writes); - } - (false, true) => { - // We have to grow here because the new bits are going to get flipped to 1. - self.component_read_and_writes.grow( - self.component_read_and_writes - .len() - .max(other.component_read_and_writes.len()), - ); - self.component_read_and_writes.toggle_range(..); - self.component_read_and_writes - .intersect_with(&other.component_read_and_writes); - } - (false, false) => { - self.component_read_and_writes - .union_with(&other.component_read_and_writes); - } - } - - match ( - self.component_writes_inverted, + ); + invertible_union_with( + &mut self.component_writes, + &mut self.component_writes_inverted, + &other.component_writes, other.component_writes_inverted, - ) { - (true, true) => { - self.component_writes - .intersect_with(&other.component_writes); - } - (true, false) => { - self.component_writes - .difference_with(&other.component_writes); - } - (false, true) => { - // We have to grow here because the new bits are going to get flipped to 1. - self.component_writes.grow( - self.component_writes - .len() - .max(other.component_writes.len()), - ); - self.component_writes.toggle_range(..); - self.component_writes - .intersect_with(&other.component_writes); - } - (false, false) => { - self.component_writes.union_with(&other.component_writes); - } - } + ); self.reads_all_resources = self.reads_all_resources || other.reads_all_resources; self.writes_all_resources = self.writes_all_resources || other.writes_all_resources; - self.component_read_and_writes_inverted = component_read_and_writes_inverted; - self.component_writes_inverted = component_writes_inverted; self.resource_read_and_writes .union_with(&other.resource_read_and_writes); self.resource_writes.union_with(&other.resource_writes); self.archetypal.union_with(&other.archetypal); } + /// Removes any access from `self` that would conflict with `other`. + /// This removes any reads and writes for any component written by `other`, + /// and removes any writes for any component read by `other`. + pub fn remove_conflicting_access(&mut self, other: &Access) { + invertible_difference_with( + &mut self.component_read_and_writes, + &mut self.component_read_and_writes_inverted, + &other.component_writes, + other.component_writes_inverted, + ); + invertible_difference_with( + &mut self.component_writes, + &mut self.component_writes_inverted, + &other.component_read_and_writes, + other.component_read_and_writes_inverted, + ); + + if other.reads_all_resources { + self.writes_all_resources = false; + self.resource_writes.clear(); + } + if other.writes_all_resources { + self.reads_all_resources = false; + self.resource_read_and_writes.clear(); + } + self.resource_read_and_writes + .difference_with(&other.resource_writes); + self.resource_writes + .difference_with(&other.resource_read_and_writes); + } + /// Returns `true` if the access and `other` can be active at the same time, /// only looking at their component access. /// @@ -840,6 +820,55 @@ impl Access { } } +/// Performs an in-place union of `other` into `self`, where either set may be inverted. +/// +/// Each set corresponds to a `FixedBitSet` if `inverted` is `false`, +/// or to the infinite (co-finite) complement of the `FixedBitSet` if `inverted` is `true`. +/// +/// This updates the `self` set to include any elements in the `other` set. +/// Note that this may change `self_inverted` to `true` if we add an infinite +/// set to a finite one, resulting in a new infinite set. +fn invertible_union_with( + self_set: &mut FixedBitSet, + self_inverted: &mut bool, + other_set: &FixedBitSet, + other_inverted: bool, +) { + match (*self_inverted, other_inverted) { + (true, true) => self_set.intersect_with(other_set), + (true, false) => self_set.difference_with(other_set), + (false, true) => { + *self_inverted = true; + // We have to grow here because the new bits are going to get flipped to 1. + self_set.grow(other_set.len()); + self_set.toggle_range(..); + self_set.intersect_with(other_set); + } + (false, false) => self_set.union_with(other_set), + } +} + +/// Performs an in-place set difference of `other` from `self`, where either set may be inverted. +/// +/// Each set corresponds to a `FixedBitSet` if `inverted` is `false`, +/// or to the infinite (co-finite) complement of the `FixedBitSet` if `inverted` is `true`. +/// +/// This updates the `self` set to remove any elements in the `other` set. +/// Note that this may change `self_inverted` to `false` if we remove an +/// infinite set from another infinite one, resulting in a finite difference. +fn invertible_difference_with( + self_set: &mut FixedBitSet, + self_inverted: &mut bool, + other_set: &FixedBitSet, + other_inverted: bool, +) { + // We can share the implementation of `invertible_union_with` with some algebra: + // A - B = A & !B = !(!A | B) + *self_inverted = !*self_inverted; + invertible_union_with(self_set, self_inverted, other_set, other_inverted); + *self_inverted = !*self_inverted; +} + /// Error returned when attempting to iterate over items included in an [`Access`] /// if the access excludes items rather than including them. #[derive(Clone, Copy, PartialEq, Eq, Debug, Error)] @@ -1428,6 +1457,7 @@ impl Default for FilteredAccessSet { #[cfg(test)] mod tests { + use super::{invertible_difference_with, invertible_union_with}; use crate::query::{ access::AccessFilters, Access, AccessConflicts, ComponentAccessKind, FilteredAccess, FilteredAccessSet, UnboundedAccessError, @@ -1770,4 +1800,99 @@ mod tests { }), ); } + + /// Create a `FixedBitSet` with a given number of total bits and a given list of bits to set. + /// Setting the number of bits is important in tests since the `PartialEq` impl checks that the length matches. + fn bit_set(bits: usize, iter: impl IntoIterator) -> FixedBitSet { + let mut result = FixedBitSet::with_capacity(bits); + result.extend(iter); + result + } + + #[test] + fn invertible_union_with_tests() { + let invertible_union = |mut self_inverted: bool, other_inverted: bool| { + // Check all four possible bit states: In both sets, the first, the second, or neither + let mut self_set = bit_set(4, [0, 1]); + let other_set = bit_set(4, [0, 2]); + invertible_union_with( + &mut self_set, + &mut self_inverted, + &other_set, + other_inverted, + ); + (self_set, self_inverted) + }; + + // Check each combination of `inverted` flags + let (s, i) = invertible_union(false, false); + // [0, 1] | [0, 2] = [0, 1, 2] + assert_eq!((s, i), (bit_set(4, [0, 1, 2]), false)); + + let (s, i) = invertible_union(false, true); + // [0, 1] | [1, 3, ...] = [0, 1, 3, ...] + assert_eq!((s, i), (bit_set(4, [2]), true)); + + let (s, i) = invertible_union(true, false); + // [2, 3, ...] | [0, 2] = [0, 2, 3, ...] + assert_eq!((s, i), (bit_set(4, [1]), true)); + + let (s, i) = invertible_union(true, true); + // [2, 3, ...] | [1, 3, ...] = [1, 2, 3, ...] + assert_eq!((s, i), (bit_set(4, [0]), true)); + } + + #[test] + fn invertible_union_with_different_lengths() { + // When adding a large inverted set to a small normal set, + // make sure we invert the bits beyond the original length. + // Failing to call `grow` before `toggle_range` would cause bit 1 to be zero, + // which would incorrectly treat it as included in the output set. + let mut self_set = bit_set(1, [0]); + let mut self_inverted = false; + let other_set = bit_set(3, [0, 1]); + let other_inverted = true; + invertible_union_with( + &mut self_set, + &mut self_inverted, + &other_set, + other_inverted, + ); + + // [0] | [2, ...] = [0, 2, ...] + assert_eq!((self_set, self_inverted), (bit_set(3, [1]), true)); + } + + #[test] + fn invertible_difference_with_tests() { + let invertible_difference = |mut self_inverted: bool, other_inverted: bool| { + // Check all four possible bit states: In both sets, the first, the second, or neither + let mut self_set = bit_set(4, [0, 1]); + let other_set = bit_set(4, [0, 2]); + invertible_difference_with( + &mut self_set, + &mut self_inverted, + &other_set, + other_inverted, + ); + (self_set, self_inverted) + }; + + // Check each combination of `inverted` flags + let (s, i) = invertible_difference(false, false); + // [0, 1] - [0, 2] = [1] + assert_eq!((s, i), (bit_set(4, [1]), false)); + + let (s, i) = invertible_difference(false, true); + // [0, 1] - [1, 3, ...] = [0] + assert_eq!((s, i), (bit_set(4, [0]), false)); + + let (s, i) = invertible_difference(true, false); + // [2, 3, ...] - [0, 2] = [3, ...] + assert_eq!((s, i), (bit_set(4, [0, 1, 2]), true)); + + let (s, i) = invertible_difference(true, true); + // [2, 3, ...] - [1, 3, ...] = [2] + assert_eq!((s, i), (bit_set(4, [2]), false)); + } } diff --git a/crates/bevy_ecs/src/query/builder.rs b/crates/bevy_ecs/src/query/builder.rs index 81819cb9ac..b545caad8f 100644 --- a/crates/bevy_ecs/src/query/builder.rs +++ b/crates/bevy_ecs/src/query/builder.rs @@ -248,11 +248,9 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> { pub fn transmute_filtered( &mut self, ) -> &mut QueryBuilder<'w, NewD, NewF> { - let mut fetch_state = NewD::init_state(self.world); + let fetch_state = NewD::init_state(self.world); let filter_state = NewF::init_state(self.world); - NewD::set_access(&mut fetch_state, &self.access); - let mut access = FilteredAccess::default(); NewD::update_component_access(&fetch_state, &mut access); NewF::update_component_access(&filter_state, &mut access); @@ -275,7 +273,10 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> { #[cfg(test)] mod tests { - use crate::{prelude::*, world::FilteredEntityRef}; + use crate::{ + prelude::*, + world::{EntityMutExcept, EntityRefExcept, FilteredEntityMut, FilteredEntityRef}, + }; use std::dbg; #[derive(Component, PartialEq, Debug)] @@ -422,6 +423,89 @@ mod tests { } } + #[test] + fn builder_provide_access() { + let mut world = World::new(); + world.spawn((A(0), B(1))); + + let mut query = + QueryBuilder::<(Entity, FilteredEntityRef, FilteredEntityMut)>::new(&mut world) + .data::<&mut A>() + .data::<&B>() + .build(); + + // The `FilteredEntityRef` only has read access, so the `FilteredEntityMut` can have read access without conflicts + let (_entity, entity_ref_1, mut entity_ref_2) = query.single_mut(&mut world).unwrap(); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_2.get::().is_some()); + assert!(entity_ref_2.get_mut::().is_none()); + assert!(entity_ref_2.get::().is_some()); + assert!(entity_ref_2.get_mut::().is_none()); + + let mut query = + QueryBuilder::<(Entity, FilteredEntityMut, FilteredEntityMut)>::new(&mut world) + .data::<&mut A>() + .data::<&B>() + .build(); + + // The first `FilteredEntityMut` has write access to A, so the second one cannot have write access + let (_entity, mut entity_ref_1, mut entity_ref_2) = query.single_mut(&mut world).unwrap(); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_1.get_mut::().is_some()); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_1.get_mut::().is_none()); + assert!(entity_ref_2.get::().is_none()); + assert!(entity_ref_2.get_mut::().is_none()); + assert!(entity_ref_2.get::().is_some()); + assert!(entity_ref_2.get_mut::().is_none()); + + let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B)>::new(&mut world) + .data::<&mut A>() + .data::<&mut B>() + .build(); + + // Any `A` access would conflict with `&mut A`, and write access to `B` would conflict with `&B`. + let (mut entity_ref, _a, _b) = query.single_mut(&mut world).unwrap(); + assert!(entity_ref.get::().is_none()); + assert!(entity_ref.get_mut::().is_none()); + assert!(entity_ref.get::().is_some()); + assert!(entity_ref.get_mut::().is_none()); + + let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B)>::new(&mut world) + .data::() + .build(); + + // Same as above, but starting from "all" access + let (mut entity_ref, _a, _b) = query.single_mut(&mut world).unwrap(); + assert!(entity_ref.get::().is_none()); + assert!(entity_ref.get_mut::().is_none()); + assert!(entity_ref.get::().is_some()); + assert!(entity_ref.get_mut::().is_none()); + + let mut query = QueryBuilder::<(FilteredEntityMut, EntityMutExcept)>::new(&mut world) + .data::() + .build(); + + // Removing `EntityMutExcept` just leaves A + let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap(); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_1.get_mut::().is_some()); + assert!(entity_ref_1.get::().is_none()); + assert!(entity_ref_1.get_mut::().is_none()); + + let mut query = QueryBuilder::<(FilteredEntityMut, EntityRefExcept)>::new(&mut world) + .data::() + .build(); + + // Removing `EntityRefExcept` just leaves A, plus read access + let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap(); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_1.get_mut::().is_some()); + assert!(entity_ref_1.get::().is_some()); + assert!(entity_ref_1.get_mut::().is_none()); + } + /// Regression test for issue #14348 #[test] fn builder_static_dense_dynamic_sparse() { diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index cd632f7b14..c244f7fbcb 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -291,6 +291,22 @@ pub unsafe trait QueryData: WorldQuery { /// This function manually implements subtyping for the query items. fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort>; + /// Offers additional access above what we requested in `update_component_access`. + /// Implementations may add additional access that is a subset of `available_access` + /// and does not conflict with anything in `access`, + /// and must update `access` to include that access. + /// + /// This is used by [`WorldQuery`] types like [`FilteredEntityRef`] + /// and [`FilteredEntityMut`] to support dynamic access. + /// + /// Called when constructing a [`QueryLens`](crate::system::QueryLens) or calling [`QueryState::from_builder`](super::QueryState::from_builder) + fn provide_extra_access( + _state: &mut Self::State, + _access: &mut Access, + _available_access: &Access, + ) { + } + /// Fetch [`Self::Item`](`QueryData::Item`) for either the given `entity` in the current [`Table`], /// or for the given `entity` in the current [`Archetype`]. This must always be called after /// [`WorldQuery::set_table`] with a `table_row` in the range of the current [`Table`] or after @@ -635,7 +651,7 @@ unsafe impl<'a> QueryData for EntityMut<'a> { /// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self` unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> { type Fetch<'w> = (UnsafeWorldCell<'w>, Access); - type State = FilteredAccess; + type State = Access; fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { fetch @@ -661,18 +677,12 @@ unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> { _: &'w Archetype, _table: &Table, ) { - fetch.1.clone_from(&state.access); + fetch.1.clone_from(state); } #[inline] unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, _: &'w Table) { - fetch.1.clone_from(&state.access); - } - - #[inline] - fn set_access<'w>(state: &mut Self::State, access: &FilteredAccess) { - state.clone_from(access); - state.access_mut().clear_writes(); + fetch.1.clone_from(state); } fn update_component_access( @@ -680,18 +690,18 @@ unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> { filtered_access: &mut FilteredAccess, ) { assert!( - filtered_access.access().is_compatible(&state.access), + filtered_access.access().is_compatible(state), "FilteredEntityRef conflicts with a previous access in this query. Exclusive access cannot coincide with any other accesses.", ); - filtered_access.access.extend(&state.access); + filtered_access.access.extend(state); } fn init_state(_world: &mut World) -> Self::State { - FilteredAccess::default() + Access::default() } fn get_state(_components: &Components) -> Option { - Some(FilteredAccess::default()) + Some(Access::default()) } fn matches_component_set( @@ -712,6 +722,25 @@ unsafe impl<'a> QueryData for FilteredEntityRef<'a> { item } + #[inline] + fn provide_extra_access( + state: &mut Self::State, + access: &mut Access, + available_access: &Access, + ) { + // Claim any extra access that doesn't conflict with other subqueries + // This is used when constructing a `QueryLens` or creating a query from a `QueryBuilder` + // Start with the entire available access, since that is the most we can possibly access + state.clone_from(available_access); + // Prevent all writes, since `FilteredEntityRef` only performs read access + state.clear_writes(); + // Prevent any access that would conflict with other accesses in the current query + state.remove_conflicting_access(access); + // Finally, add the resulting access to the query access + // to make sure a later `FilteredEntityMut` won't conflict with this. + access.extend(state); + } + #[inline(always)] unsafe fn fetch<'w>( (world, access): &mut Self::Fetch<'w>, @@ -731,7 +760,7 @@ unsafe impl ReadOnlyQueryData for FilteredEntityRef<'_> {} /// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self` unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> { type Fetch<'w> = (UnsafeWorldCell<'w>, Access); - type State = FilteredAccess; + type State = Access; fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { fetch @@ -757,17 +786,12 @@ unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> { _: &'w Archetype, _table: &Table, ) { - fetch.1.clone_from(&state.access); + fetch.1.clone_from(state); } #[inline] unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, _: &'w Table) { - fetch.1.clone_from(&state.access); - } - - #[inline] - fn set_access<'w>(state: &mut Self::State, access: &FilteredAccess) { - state.clone_from(access); + fetch.1.clone_from(state); } fn update_component_access( @@ -775,18 +799,18 @@ unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> { filtered_access: &mut FilteredAccess, ) { assert!( - filtered_access.access().is_compatible(&state.access), + filtered_access.access().is_compatible(state), "FilteredEntityMut conflicts with a previous access in this query. Exclusive access cannot coincide with any other accesses.", ); - filtered_access.access.extend(&state.access); + filtered_access.access.extend(state); } fn init_state(_world: &mut World) -> Self::State { - FilteredAccess::default() + Access::default() } fn get_state(_components: &Components) -> Option { - Some(FilteredAccess::default()) + Some(Access::default()) } fn matches_component_set( @@ -807,6 +831,23 @@ unsafe impl<'a> QueryData for FilteredEntityMut<'a> { item } + #[inline] + fn provide_extra_access( + state: &mut Self::State, + access: &mut Access, + available_access: &Access, + ) { + // Claim any extra access that doesn't conflict with other subqueries + // This is used when constructing a `QueryLens` or creating a query from a `QueryBuilder` + // Start with the entire available access, since that is the most we can possibly access + state.clone_from(available_access); + // Prevent any access that would conflict with other accesses in the current query + state.remove_conflicting_access(access); + // Finally, add the resulting access to the query access + // to make sure a later `FilteredEntityRef` or `FilteredEntityMut` won't conflict with this. + access.extend(state); + } + #[inline(always)] unsafe fn fetch<'w>( (world, access): &mut Self::Fetch<'w>, @@ -2079,6 +2120,16 @@ macro_rules! impl_tuple_query_data { )*) } + #[inline] + fn provide_extra_access( + state: &mut Self::State, + access: &mut Access, + available_access: &Access, + ) { + let ($($name,)*) = state; + $($name::provide_extra_access($name, access, available_access);)* + } + #[inline(always)] unsafe fn fetch<'w>( fetch: &mut Self::Fetch<'w>, diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 6c2d26b08c..0591bac58c 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -287,7 +287,14 @@ impl QueryState { pub fn from_builder(builder: &mut QueryBuilder) -> Self { let mut fetch_state = D::init_state(builder.world_mut()); let filter_state = F::init_state(builder.world_mut()); - D::set_access(&mut fetch_state, builder.access()); + + let mut component_access = FilteredAccess::default(); + D::update_component_access(&fetch_state, &mut component_access); + D::provide_extra_access( + &mut fetch_state, + component_access.access_mut(), + builder.access().access(), + ); let mut component_access = builder.access().clone(); @@ -753,29 +760,27 @@ impl QueryState { let mut fetch_state = NewD::get_state(world.components()).expect("Could not create fetch_state, Please initialize all referenced components before transmuting."); let filter_state = NewF::get_state(world.components()).expect("Could not create filter_state, Please initialize all referenced components before transmuting."); - fn to_readonly(mut access: FilteredAccess) -> FilteredAccess { - access.access_mut().clear_writes(); - access - } - - let self_access = if D::IS_READ_ONLY && self.component_access.access().has_any_write() { + let mut self_access = self.component_access.clone(); + if D::IS_READ_ONLY { // The current state was transmuted from a mutable // `QueryData` to a read-only one. // Ignore any write access in the current state. - &to_readonly(self.component_access.clone()) - } else { - &self.component_access - }; + self_access.access_mut().clear_writes(); + } - NewD::set_access(&mut fetch_state, self_access); NewD::update_component_access(&fetch_state, &mut component_access); + NewD::provide_extra_access( + &mut fetch_state, + component_access.access_mut(), + self_access.access(), + ); let mut filter_component_access = FilteredAccess::default(); NewF::update_component_access(&filter_state, &mut filter_component_access); component_access.extend(&filter_component_access); assert!( - component_access.is_subset(self_access), + component_access.is_subset(&self_access), "Transmuted state for {} attempts to access terms that are not allowed by original state {}.", core::any::type_name::<(NewD, NewF)>(), core::any::type_name::<(D, F)>() ); @@ -787,7 +792,7 @@ impl QueryState { is_dense: self.is_dense, fetch_state, filter_state, - component_access: self.component_access.clone(), + component_access: self_access, matched_tables: self.matched_tables.clone(), matched_archetypes: self.matched_archetypes.clone(), #[cfg(feature = "trace")] @@ -881,8 +886,12 @@ impl QueryState { } } - NewD::set_access(&mut new_fetch_state, &joined_component_access); NewD::update_component_access(&new_fetch_state, &mut component_access); + NewD::provide_extra_access( + &mut new_fetch_state, + component_access.access_mut(), + joined_component_access.access(), + ); let mut new_filter_component_access = FilteredAccess::default(); NewF::update_component_access(&new_filter_state, &mut new_filter_component_access); @@ -2062,12 +2071,12 @@ mod tests { fn can_transmute_filtered_entity() { let mut world = World::new(); let entity = world.spawn((A(0), B(1))).id(); - let query = - QueryState::<(Entity, &A, &B)>::new(&mut world).transmute::(&world); + let query = QueryState::<(Entity, &A, &B)>::new(&mut world) + .transmute::<(Entity, FilteredEntityRef)>(&world); let mut query = query; // Our result is completely untyped - let entity_ref = query.single(&world).unwrap(); + let (_entity, entity_ref) = query.single(&world).unwrap(); assert_eq!(entity, entity_ref.id()); assert_eq!(0, entity_ref.get::().unwrap().0); @@ -2285,11 +2294,11 @@ mod tests { let query_1 = QueryState::<&mut A>::new(&mut world); let query_2 = QueryState::<&mut B>::new(&mut world); - let mut new_query: QueryState = query_1.join(&world, &query_2); + let mut new_query: QueryState<(Entity, FilteredEntityMut)> = query_1.join(&world, &query_2); - let mut entity = new_query.single_mut(&mut world).unwrap(); - assert!(entity.get_mut::().is_some()); - assert!(entity.get_mut::().is_some()); + let (_entity, mut entity_mut) = new_query.single_mut(&mut world).unwrap(); + assert!(entity_mut.get_mut::().is_some()); + assert!(entity_mut.get_mut::().is_some()); } #[test] diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index da147770e0..a6bcbf58bd 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/crates/bevy_ecs/src/query/world_query.rs @@ -13,11 +13,11 @@ use variadics_please::all_tuples; /// # Safety /// /// Implementor must ensure that -/// [`update_component_access`], [`matches_component_set`], [`QueryData::fetch`], [`QueryFilter::filter_fetch`] and [`init_fetch`] +/// [`update_component_access`], [`QueryData::provide_extra_access`], [`matches_component_set`], [`QueryData::fetch`], [`QueryFilter::filter_fetch`] and [`init_fetch`] /// obey the following: /// -/// - For each component mutably accessed by [`QueryData::fetch`], [`update_component_access`] should add write access unless read or write access has already been added, in which case it should panic. -/// - For each component readonly accessed by [`QueryData::fetch`] or [`QueryFilter::filter_fetch`], [`update_component_access`] should add read access unless write access has already been added, in which case it should panic. +/// - For each component mutably accessed by [`QueryData::fetch`], [`update_component_access`] or [`QueryData::provide_extra_access`] should add write access unless read or write access has already been added, in which case it should panic. +/// - For each component readonly accessed by [`QueryData::fetch`] or [`QueryFilter::filter_fetch`], [`update_component_access`] or [`QueryData::provide_extra_access`] should add read access unless write access has already been added, in which case it should panic. /// - If `fetch` mutably accesses the same component twice, [`update_component_access`] should panic. /// - [`update_component_access`] may not add a `Without` filter for a component unless [`matches_component_set`] always returns `false` when the component set contains that component. /// - [`update_component_access`] may not add a `With` filter for a component unless [`matches_component_set`] always returns `false` when the component set doesn't contain that component. @@ -27,9 +27,11 @@ use variadics_please::all_tuples; /// - Each filter in that disjunction must be a conjunction of the corresponding element's filter with the previous `access` /// - For each resource readonly accessed by [`init_fetch`], [`update_component_access`] should add read access. /// - Mutable resource access is not allowed. +/// - Any access added during [`QueryData::provide_extra_access`] must be a subset of `available_access`, and must not conflict with any access in `access`. /// /// When implementing [`update_component_access`], note that `add_read` and `add_write` both also add a `With` filter, whereas `extend_access` does not change the filters. /// +/// [`QueryData::provide_extra_access`]: crate::query::QueryData::provide_extra_access /// [`QueryData::fetch`]: crate::query::QueryData::fetch /// [`QueryFilter::filter_fetch`]: crate::query::QueryFilter::filter_fetch /// [`init_fetch`]: Self::init_fetch @@ -101,12 +103,6 @@ pub unsafe trait WorldQuery { /// - `state` must be the [`State`](Self::State) that `fetch` was initialized with. unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table); - /// Sets available accesses for implementors with dynamic access such as [`FilteredEntityRef`](crate::world::FilteredEntityRef) - /// or [`FilteredEntityMut`](crate::world::FilteredEntityMut). - /// - /// Called when constructing a [`QueryLens`](crate::system::QueryLens) or calling [`QueryState::from_builder`](super::QueryState::from_builder) - fn set_access(_state: &mut Self::State, _access: &FilteredAccess) {} - /// Adds any component accesses used by this [`WorldQuery`] to `access`. /// /// Used to check which queries are disjoint and can run in parallel diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 183bdecfb4..627a02f6ea 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -2208,8 +2208,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// so can be added to any query /// * [`FilteredEntityRef`](crate::world::FilteredEntityRef) and [`FilteredEntityMut`](crate::world::FilteredEntityMut) /// have access determined by the [`QueryBuilder`](crate::query::QueryBuilder) used to construct them. - /// Any query can be transmuted to them, and they will receive the access of the source query, - /// but only if they are the top-level query and not nested + /// Any query can be transmuted to them, and they will receive the access of the source query. + /// When combined with other `QueryData`, they will receive any access of the source query that does not conflict with the other data. /// * [`Added`](crate::query::Added) and [`Changed`](crate::query::Changed) filters have read and required access to `T` /// * [`With`](crate::query::With) and [`Without`](crate::query::Without) filters have no access at all, /// so can be added to any query @@ -2281,9 +2281,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// // Anything can be transmuted to `FilteredEntityRef` or `FilteredEntityMut` /// // This will create a `FilteredEntityMut` that only has read access to `T` /// assert_valid_transmute::<&T, FilteredEntityMut>(); - /// // This transmute will succeed, but the `FilteredEntityMut` will have no access! - /// // It must be the top-level query to be given access, but here it is nested in a tuple. - /// assert_valid_transmute::<&T, (Entity, FilteredEntityMut)>(); + /// // This will create a `FilteredEntityMut` that has no access to `T`, + /// // read access to `U`, and write access to `V`. + /// assert_valid_transmute::<(&mut T, &mut U, &mut V), (&mut T, &U, FilteredEntityMut)>(); /// /// // `Added` and `Changed` filters have the same access as `&T` data /// // Remember that they are only evaluated on the transmuted query, not the original query! diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index a0ef602cfc..2346767f2f 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -3346,14 +3346,6 @@ impl<'w, 'a, T: Component> VacantEntry<'w, 'a, T> { /// /// let filtered_entity: FilteredEntityRef = query.single(&mut world).unwrap(); /// let component: &A = filtered_entity.get().unwrap(); -/// -/// // Here `FilteredEntityRef` is nested in a tuple, so it does not have access to `&A`. -/// let mut query = QueryBuilder::<(Entity, FilteredEntityRef)>::new(&mut world) -/// .data::<&A>() -/// .build(); -/// -/// let (_, filtered_entity) = query.single(&mut world).unwrap(); -/// assert!(filtered_entity.get::().is_none()); /// ``` #[derive(Clone)] pub struct FilteredEntityRef<'w> { @@ -3677,14 +3669,6 @@ unsafe impl EntityEquivalent for FilteredEntityRef<'_> {} /// /// let mut filtered_entity: FilteredEntityMut = query.single_mut(&mut world).unwrap(); /// let component: Mut = filtered_entity.get_mut().unwrap(); -/// -/// // Here `FilteredEntityMut` is nested in a tuple, so it does not have access to `&mut A`. -/// let mut query = QueryBuilder::<(Entity, FilteredEntityMut)>::new(&mut world) -/// .data::<&mut A>() -/// .build(); -/// -/// let (_, mut filtered_entity) = query.single_mut(&mut world).unwrap(); -/// assert!(filtered_entity.get_mut::().is_none()); /// ``` pub struct FilteredEntityMut<'w> { entity: UnsafeEntityCell<'w>,