diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 13cd5501e5a81..7ae402536806a 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -3049,11 +3049,10 @@ pub struct NestedQueryFetch<'w> { // SAFETY: // Does not access any components on the current entity // Accesses through the nested query are registered in `init_nested_access` -unsafe impl WorldQuery - for NestedQuery -{ +unsafe impl WorldQuery for NestedQuery { type Fetch<'w> = NestedQueryFetch<'w>; - type State = QueryState; + // Note: The state's QueryData should be treated as D via as_transmuted_state(_mut) + type State = QueryState; fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { fetch @@ -3099,6 +3098,9 @@ unsafe impl WorldQuery component_access_set: &mut FilteredAccessSet, world: UnsafeWorldCell, ) { + // SAFETY: D is the original QueryData for the QueryState. + let state = unsafe { state.as_transmuted_state::() }; + state.init_access(system_name, component_access_set, world); } @@ -3107,7 +3109,7 @@ unsafe impl WorldQuery // `WorldQuery::init_nested_access` must be called before `WorldQuery::init_fetch, // which must be called before `QueryData::fetch`, // and we only call methods on the `QueryState` in `fetch`. - unsafe { QueryState::::new_unchecked(world) } + unsafe { QueryState::::new_unchecked(world) }.to_readonly() } fn get_state(_components: &Components) -> Option { @@ -3126,6 +3128,9 @@ unsafe impl WorldQuery } fn update_archetypes(state: &mut Self::State, world: UnsafeWorldCell) { + // SAFETY: D is the original QueryData for the QueryState. + let state = unsafe { state.as_transmuted_state_mut::() }; + state.update_archetypes_unsafe_world_cell(world); } } @@ -3133,16 +3138,14 @@ unsafe impl WorldQuery // SAFETY: // `Self::ReadOnly` accesses `D::ReadOnly`, which is a subset of the data accessed by `D` // `IS_READ_ONLY` iff `D::IS_READ_ONLY` iff `D: ReadOnlyQueryData` iff `Self: ReadOnlyQueryData` -unsafe impl QueryData - for NestedQuery -{ +unsafe impl QueryData for NestedQuery { const IS_READ_ONLY: bool = D::IS_READ_ONLY; // Nested queries are always archetypal because `fetch` always returns `Some`. // If `D::IS_ARCHETYPAL == false` or `F::IS_ARCHETYPAL == false`, // then the nested query may filter out some entities that *it* matches, // but it will not filter the outer query. const IS_ARCHETYPAL: bool = true; - type ReadOnly = NestedQuery; + type ReadOnly = NestedQuery; type Item<'w, 's> = Query<'w, 's, D, F>; fn shrink<'wlong: 'wshort, 'wshort, 's>( @@ -3158,6 +3161,9 @@ unsafe impl QueryData _entity: Entity, _table_row: TableRow, ) -> Option> { + // SAFETY: D is the original QueryData for the QueryState. + let state = unsafe { state.as_transmuted_state::() }; + // SAFETY: // - We registered the required access in `init_nested_access`, so it's available. // - If we are fetching multiple entities concurrently, diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 083d47ca0f1ce..ae2dfc99fbff5 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -120,6 +120,37 @@ impl FromWorld for QueryState { } impl QueryState { + /// Converts this `QueryState` to a `QueryState` that does not access anything mutably. + pub fn to_readonly(self) -> QueryState { + let QueryState { + world_id, + archetype_generation, + matched_tables, + matched_archetypes, + component_access, + matched_storage_ids, + is_dense, + fetch_state, + filter_state, + #[cfg(feature = "trace")] + par_iter_span, + } = self; + + QueryState { + world_id, + archetype_generation, + matched_tables, + matched_archetypes, + component_access, + matched_storage_ids, + is_dense, + fetch_state, + filter_state, + #[cfg(feature = "trace")] + par_iter_span, + } + } + /// Converts this `QueryState` reference to a `QueryState` that does not access anything mutably. pub fn as_readonly(&self) -> &QueryState { // SAFETY: invariant on `WorldQuery` trait upholds that `D::ReadOnly` and `F::ReadOnly` @@ -148,7 +179,7 @@ impl QueryState { /// `NewD` must have a subset of the access that `D` does and match the exact same archetypes/tables /// `NewF` must have a subset of the access that `F` does and match the exact same archetypes/tables pub(crate) unsafe fn as_transmuted_state< - NewD: ReadOnlyQueryData, + NewD: QueryData, NewF: QueryFilter, >( &self, @@ -156,6 +187,22 @@ impl QueryState { &*ptr::from_ref(self).cast::>() } + /// Converts this `QueryState` reference to any other `QueryState` with + /// the same `WorldQuery::State` associated types. + /// + /// # Safety + /// + /// `NewD` must have a subset of the access that `D` does and match the exact same archetypes/tables + /// `NewF` must have a subset of the access that `F` does and match the exact same archetypes/tables + pub(crate) unsafe fn as_transmuted_state_mut< + NewD: QueryData, + NewF: QueryFilter, + >( + &mut self, + ) -> &mut QueryState { + &mut *ptr::from_mut(self).cast::>() + } + /// Returns the components accessed by this query. pub fn component_access(&self) -> &FilteredAccess { &self.component_access diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index 11328d2223c69..5d05bd66be91f 100644 --- a/crates/bevy_ecs/src/system/mod.rs +++ b/crates/bevy_ecs/src/system/mod.rs @@ -918,6 +918,33 @@ mod tests { run_system(&mut world, sys); } + #[test] + #[should_panic = "error[B0001]"] + fn mut_nested_query_conflicts_with_main_query() { + fn sys(_: Query<(&A, NestedQuery<&mut A>)>) {} + + let mut world = World::default(); + run_system(&mut world, sys); + } + + #[test] + #[should_panic = "error[B0001]"] + fn mut_nested_query_conflicts_with_earlier_query() { + fn sys(_: Query<&A>, _: Query>) {} + + let mut world = World::default(); + run_system(&mut world, sys); + } + + #[test] + #[should_panic = "error[B0001]"] + fn mut_nested_query_conflicts_with_later_query() { + fn sys(_: Query>, _: Query<&A>) {} + + let mut world = World::default(); + run_system(&mut world, sys); + } + #[test] fn query_set_system() { fn sys(mut _set: ParamSet<(Query<&mut A>, Query<&A>)>) {}