Allow write access in NestedQuery

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<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> WorldQuery
-    for NestedQuery<D, F>
-{
+unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> WorldQuery for NestedQuery<D, F> {
     type Fetch<'w> = NestedQueryFetch<'w>;
-    type State = QueryState<D, F>;
+    // Note: The state's QueryData should be treated as D via as_transmuted_state(_mut)
+    type State = QueryState<D::ReadOnly, F>;
 
     fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
         fetch
@@ -3099,6 +3098,9 @@ unsafe impl<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> WorldQuery
         component_access_set: &mut FilteredAccessSet,
         world: UnsafeWorldCell,
     ) {
+        // SAFETY: D is the original QueryData for the QueryState.
+        let state = unsafe { state.as_transmuted_state::<D, F>() };
+
         state.init_access(system_name, component_access_set, world);
     }
 
@@ -3107,7 +3109,7 @@ unsafe impl<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> 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::<D, F>::new_unchecked(world) }
+        unsafe { QueryState::<D, F>::new_unchecked(world) }.to_readonly()
     }
 
     fn get_state(_components: &Components) -> Option<Self::State> {
@@ -3126,23 +3128,24 @@ unsafe impl<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> 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::<D, F>() };
+
         state.update_archetypes_unsafe_world_cell(world);
     }
 }
 
 // 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<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> QueryData
-    for NestedQuery<D, F>
-{
+unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> QueryData for NestedQuery<D, F> {
     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<D, F>;
+    type ReadOnly = NestedQuery<D::ReadOnly, F>;
     type Item<'w, 's> = Query<'w, 's, D, F>;
 
     fn shrink<'wlong: 'wshort, 'wshort, 's>(
@@ -3158,6 +3161,9 @@ unsafe impl<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> QueryData
         _entity: Entity,
         _table_row: TableRow,
     ) -> Option<Self::Item<'w, 's>> {
+        // SAFETY: D is the original QueryData for the QueryState.
+        let state = unsafe { state.as_transmuted_state::<D, F>() };
+
         // SAFETY:
         // - We registered the required access in `init_nested_access`, so it's available.
         // - If we are fetching multiple entities concurrently,

crates/bevy_ecs/src/query/state.rs

@@ -120,6 +120,37 @@ impl<D: QueryData, F: QueryFilter> FromWorld for QueryState<D, F> {
 }
 
 impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
+    /// Converts this `QueryState` to a `QueryState` that does not access anything mutably.
+    pub fn to_readonly(self) -> QueryState<D::ReadOnly, F> {
+        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<D::ReadOnly, F> {
         // SAFETY: invariant on `WorldQuery` trait upholds that `D::ReadOnly` and `F::ReadOnly`
@@ -148,14 +179,30 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     /// `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<State = D::State>,
+        NewD: QueryData<State = D::State>,
         NewF: QueryFilter<State = F::State>,
     >(
         &self,
     ) -> &QueryState<NewD, NewF> {
         &*ptr::from_ref(self).cast::<QueryState<NewD, NewF>>()
     }
 
+    /// 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<State = D::State>,
+        NewF: QueryFilter<State = F::State>,
+    >(
+        &mut self,
+    ) -> &mut QueryState<NewD, NewF> {
+        &mut *ptr::from_mut(self).cast::<QueryState<NewD, NewF>>()
+    }
+
     /// Returns the components accessed by this query.
     pub fn component_access(&self) -> &FilteredAccess {
         &self.component_access

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<NestedQuery<&mut A>>) {}
+
+        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<NestedQuery<&mut A>>, _: 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>)>) {}