Support declaring resource access in Queries. (bevyengine#16843)

# Objective

Allow resources to be accessed soundly by `QueryData` and `QueryFilter`
implementations.

This mostly works today, and is used in `bevy-trait-query` and will be
used by bevyengine#16810. The problem is that the access is not made visible to
the executor, so it would be possible for a system with resource access
in a query to run concurrently with a system that accesses the resource
with `ResMut`, resulting in Undefined Behavior.

## Solution

Define calling `add_resource_read` or `add_resource_write` in
`WorldQuery::update_component_access` to be a supported way to declare
resource access in a query.
Modify `QueryState::new_with_access` to check for resource access and
report it in `archetype_component_acccess`.
Modify `FilteredAccess::is_compatible` to consider resource access
conflicting even on queries with disjoint filters.

crates/bevy_ecs/src/query/access.rs

@@ -1017,7 +1017,13 @@ impl<T: SparseSetIndex> FilteredAccess<T> {
 
     /// Returns `true` if this and `other` can be active at the same time.
     pub fn is_compatible(&self, other: &FilteredAccess<T>) -> bool {
-        if self.access.is_compatible(&other.access) {
+        // Resources are read from the world rather than the filtered archetypes,
+        // so they must be compatible even if the filters are disjoint.
+        if !self.access.is_resources_compatible(&other.access) {
+            return false;
+        }
+
+        if self.access.is_components_compatible(&other.access) {
             return true;
         }
 

crates/bevy_ecs/src/query/mod.rs

@@ -105,14 +105,19 @@ impl<T> DebugCheckedUnwrap for Option<T> {
 mod tests {
     use crate::{
         self as bevy_ecs,
-        component::Component,
-        prelude::{AnyOf, Changed, Entity, Or, QueryState, With, Without},
-        query::{ArchetypeFilter, Has, QueryCombinationIter, ReadOnlyQueryData},
+        archetype::Archetype,
+        component::{Component, ComponentId, Components, Tick},
+        prelude::{AnyOf, Changed, Entity, Or, QueryState, Res, ResMut, Resource, With, Without},
+        query::{
+            ArchetypeFilter, FilteredAccess, Has, QueryCombinationIter, QueryData,
+            ReadOnlyQueryData, WorldQuery,
+        },
         schedule::{IntoSystemConfigs, Schedule},
-        system::{IntoSystem, Query, System, SystemState},
-        world::World,
+        storage::{Table, TableRow},
+        system::{assert_is_system, IntoSystem, Query, System, SystemState},
+        world::{unsafe_world_cell::UnsafeWorldCell, World},
     };
-    use bevy_ecs_macros::{QueryData, QueryFilter};
+    use bevy_ecs_macros::QueryFilter;
     use core::{any::type_name, fmt::Debug, hash::Hash};
     use std::collections::HashSet;
 
@@ -792,4 +797,235 @@ mod tests {
         let values = world.query::<&B>().iter(&world).collect::<Vec<&B>>();
         assert_eq!(values, vec![&B(2)]);
     }
+
+    #[derive(Resource)]
+    struct R;
+
+    /// `QueryData` that performs read access on R to test that resource access is tracked
+    struct ReadsRData;
+
+    /// `QueryData` that performs write access on R to test that resource access is tracked
+    struct WritesRData;
+
+    /// SAFETY:
+    /// `update_component_access` adds resource read access for `R`.
+    /// `update_archetype_component_access` does nothing, as this accesses no components.
+    unsafe impl WorldQuery for ReadsRData {
+        type Item<'w> = ();
+        type Fetch<'w> = ();
+        type State = ComponentId;
+
+        fn shrink<'wlong: 'wshort, 'wshort>(_item: Self::Item<'wlong>) -> Self::Item<'wshort> {}
+
+        fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
+
+        unsafe fn init_fetch<'w>(
+            _world: UnsafeWorldCell<'w>,
+            _state: &Self::State,
+            _last_run: Tick,
+            _this_run: Tick,
+        ) -> Self::Fetch<'w> {
+        }
+
+        const IS_DENSE: bool = true;
+
+        #[inline]
+        unsafe fn set_archetype<'w>(
+            _fetch: &mut Self::Fetch<'w>,
+            _state: &Self::State,
+            _archetype: &'w Archetype,
+            _table: &Table,
+        ) {
+        }
+
+        #[inline]
+        unsafe fn set_table<'w>(
+            _fetch: &mut Self::Fetch<'w>,
+            _state: &Self::State,
+            _table: &'w Table,
+        ) {
+        }
+
+        #[inline(always)]
+        unsafe fn fetch<'w>(
+            _fetch: &mut Self::Fetch<'w>,
+            _entity: Entity,
+            _table_row: TableRow,
+        ) -> Self::Item<'w> {
+        }
+
+        fn update_component_access(
+            &component_id: &Self::State,
+            access: &mut FilteredAccess<ComponentId>,
+        ) {
+            assert!(
+                !access.access().has_resource_write(component_id),
+                "ReadsRData conflicts with a previous access in this query. Shared access cannot coincide with exclusive access."
+            );
+            access.add_resource_read(component_id);
+        }
+
+        fn init_state(world: &mut World) -> Self::State {
+            world.components.register_resource::<R>()
+        }
+
+        fn get_state(components: &Components) -> Option<Self::State> {
+            components.resource_id::<R>()
+        }
+
+        fn matches_component_set(
+            _state: &Self::State,
+            _set_contains_id: &impl Fn(ComponentId) -> bool,
+        ) -> bool {
+            true
+        }
+    }
+
+    /// SAFETY: `Self` is the same as `Self::ReadOnly`
+    unsafe impl QueryData for ReadsRData {
+        type ReadOnly = Self;
+    }
+
+    /// SAFETY: access is read only
+    unsafe impl ReadOnlyQueryData for ReadsRData {}
+
+    /// SAFETY:
+    /// `update_component_access` adds resource read access for `R`.
+    /// `update_archetype_component_access` does nothing, as this accesses no components.
+    unsafe impl WorldQuery for WritesRData {
+        type Item<'w> = ();
+        type Fetch<'w> = ();
+        type State = ComponentId;
+
+        fn shrink<'wlong: 'wshort, 'wshort>(_item: Self::Item<'wlong>) -> Self::Item<'wshort> {}
+
+        fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
+
+        unsafe fn init_fetch<'w>(
+            _world: UnsafeWorldCell<'w>,
+            _state: &Self::State,
+            _last_run: Tick,
+            _this_run: Tick,
+        ) -> Self::Fetch<'w> {
+        }
+
+        const IS_DENSE: bool = true;
+
+        #[inline]
+        unsafe fn set_archetype<'w>(
+            _fetch: &mut Self::Fetch<'w>,
+            _state: &Self::State,
+            _archetype: &'w Archetype,
+            _table: &Table,
+        ) {
+        }
+
+        #[inline]
+        unsafe fn set_table<'w>(
+            _fetch: &mut Self::Fetch<'w>,
+            _state: &Self::State,
+            _table: &'w Table,
+        ) {
+        }
+
+        #[inline(always)]
+        unsafe fn fetch<'w>(
+            _fetch: &mut Self::Fetch<'w>,
+            _entity: Entity,
+            _table_row: TableRow,
+        ) -> Self::Item<'w> {
+        }
+
+        fn update_component_access(
+            &component_id: &Self::State,
+            access: &mut FilteredAccess<ComponentId>,
+        ) {
+            assert!(
+                !access.access().has_resource_read(component_id),
+                "WritesRData conflicts with a previous access in this query. Shared access cannot coincide with exclusive access.",
+            );
+            access.add_resource_write(component_id);
+        }
+
+        fn init_state(world: &mut World) -> Self::State {
+            world.components.register_resource::<R>()
+        }
+
+        fn get_state(components: &Components) -> Option<Self::State> {
+            components.resource_id::<R>()
+        }
+
+        fn matches_component_set(
+            _state: &Self::State,
+            _set_contains_id: &impl Fn(ComponentId) -> bool,
+        ) -> bool {
+            true
+        }
+    }
+
+    /// SAFETY: `Self` is the same as `Self::ReadOnly`
+    unsafe impl QueryData for WritesRData {
+        type ReadOnly = ReadsRData;
+    }
+
+    #[test]
+    fn read_res_read_res_no_conflict() {
+        fn system(_q1: Query<ReadsRData, With<A>>, _q2: Query<ReadsRData, Without<A>>) {}
+        assert_is_system(system);
+    }
+
+    #[test]
+    #[should_panic]
+    fn read_res_write_res_conflict() {
+        fn system(_q1: Query<ReadsRData, With<A>>, _q2: Query<WritesRData, Without<A>>) {}
+        assert_is_system(system);
+    }
+
+    #[test]
+    #[should_panic]
+    fn write_res_read_res_conflict() {
+        fn system(_q1: Query<WritesRData, With<A>>, _q2: Query<ReadsRData, Without<A>>) {}
+        assert_is_system(system);
+    }
+
+    #[test]
+    #[should_panic]
+    fn write_res_write_res_conflict() {
+        fn system(_q1: Query<WritesRData, With<A>>, _q2: Query<WritesRData, Without<A>>) {}
+        assert_is_system(system);
+    }
+
+    #[test]
+    fn read_write_res_sets_archetype_component_access() {
+        let mut world = World::new();
+
+        fn read_query(_q: Query<ReadsRData, With<A>>) {}
+        let mut read_query = IntoSystem::into_system(read_query);
+        read_query.initialize(&mut world);
+
+        fn write_query(_q: Query<WritesRData, With<A>>) {}
+        let mut write_query = IntoSystem::into_system(write_query);
+        write_query.initialize(&mut world);
+
+        fn read_res(_r: Res<R>) {}
+        let mut read_res = IntoSystem::into_system(read_res);
+        read_res.initialize(&mut world);
+
+        fn write_res(_r: ResMut<R>) {}
+        let mut write_res = IntoSystem::into_system(write_res);
+        write_res.initialize(&mut world);
+
+        assert!(read_query
+            .archetype_component_access()
+            .is_compatible(read_res.archetype_component_access()));
+        assert!(!write_query
+            .archetype_component_access()
+            .is_compatible(read_res.archetype_component_access()));
+        assert!(!read_query
+            .archetype_component_access()
+            .is_compatible(write_res.archetype_component_access()));
+        assert!(!write_query
+            .archetype_component_access()
+            .is_compatible(write_res.archetype_component_access()));
+    }
 }

crates/bevy_ecs/src/query/state.rs

@@ -187,6 +187,24 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
             }
         }
         state.archetype_generation = world.archetypes.generation();
+
+        // Resource access is not part of any archetype and must be handled separately
+        if state.component_access.access().has_read_all_resources() {
+            access.read_all_resources();
+        } else {
+            for component_id in state.component_access.access().resource_reads() {
+                access.add_resource_read(world.initialize_resource_internal(component_id).id());
+            }
+        }
+
+        if state.component_access.access().has_write_all_resources() {
+            access.write_all_resources();
+        } else {
+            for component_id in state.component_access.access().resource_writes() {
+                access.add_resource_write(world.initialize_resource_internal(component_id).id());
+            }
+        }
+
         state
     }
 

crates/bevy_ecs/src/query/world_query.rs

@@ -26,10 +26,13 @@ use variadics_please::all_tuples;
 ///     - [`matches_component_set`] must be a disjunction of the element's implementations
 ///     - [`update_component_access`] must replace the filters with a disjunction of filters
 ///     - Each filter in that disjunction must be a conjunction of the corresponding element's filter with the previous `access`
+/// - For each resource mutably accessed by [`init_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 resource readonly accessed by [`init_fetch`], [`update_component_access`] should add read access unless write access has already been added, in which case it should panic.
 ///
 /// 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.
 ///
 /// [`fetch`]: Self::fetch
+/// [`init_fetch`]: Self::init_fetch
 /// [`matches_component_set`]: Self::matches_component_set
 /// [`Query`]: crate::system::Query
 /// [`update_component_access`]: Self::update_component_access