Let FilteredEntity(Ref|Mut) receive access when nested.

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>(

crates/bevy_ecs/src/query/access.rs

@@ -430,77 +430,57 @@ impl<T: SparseSetIndex> Access<T> {
 
     /// Adds all access from `other`.
     pub fn extend(&mut self, other: &Access<T>) {
-        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);
     }
 
+    /// 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<T>) {
+        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.
     ///
@@ -838,6 +818,40 @@ impl<T: SparseSetIndex> Access<T> {
     }
 }
 
+/// Performs an in-place union of `other` into `self`, where either set may be inverted.
+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.
+fn invertible_difference_with(
+    self_set: &mut FixedBitSet,
+    self_inverted: &mut bool,
+    other_set: &FixedBitSet,
+    other_inverted: bool,
+) {
+    // 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)]

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<NewD: QueryData, NewF: QueryFilter>(
         &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::<A>().is_some());
+        assert!(entity_ref_1.get::<B>().is_some());
+        assert!(entity_ref_2.get::<A>().is_some());
+        assert!(entity_ref_2.get_mut::<A>().is_none());
+        assert!(entity_ref_2.get::<B>().is_some());
+        assert!(entity_ref_2.get_mut::<B>().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::<A>().is_some());
+        assert!(entity_ref_1.get_mut::<A>().is_some());
+        assert!(entity_ref_1.get::<B>().is_some());
+        assert!(entity_ref_1.get_mut::<B>().is_none());
+        assert!(entity_ref_2.get::<A>().is_none());
+        assert!(entity_ref_2.get_mut::<A>().is_none());
+        assert!(entity_ref_2.get::<B>().is_some());
+        assert!(entity_ref_2.get_mut::<B>().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::<A>().is_none());
+        assert!(entity_ref.get_mut::<A>().is_none());
+        assert!(entity_ref.get::<B>().is_some());
+        assert!(entity_ref.get_mut::<B>().is_none());
+
+        let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B)>::new(&mut world)
+            .data::<EntityMut>()
+            .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::<A>().is_none());
+        assert!(entity_ref.get_mut::<A>().is_none());
+        assert!(entity_ref.get::<B>().is_some());
+        assert!(entity_ref.get_mut::<B>().is_none());
+
+        let mut query = QueryBuilder::<(FilteredEntityMut, EntityMutExcept<A>)>::new(&mut world)
+            .data::<EntityMut>()
+            .build();
+
+        // Removing `EntityMutExcept<A>` just leaves A
+        let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap();
+        assert!(entity_ref_1.get::<A>().is_some());
+        assert!(entity_ref_1.get_mut::<A>().is_some());
+        assert!(entity_ref_1.get::<B>().is_none());
+        assert!(entity_ref_1.get_mut::<B>().is_none());
+
+        let mut query = QueryBuilder::<(FilteredEntityMut, EntityRefExcept<A>)>::new(&mut world)
+            .data::<EntityMut>()
+            .build();
+
+        // Removing `EntityRefExcept<A>` 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::<A>().is_some());
+        assert!(entity_ref_1.get_mut::<A>().is_some());
+        assert!(entity_ref_1.get::<B>().is_some());
+        assert!(entity_ref_1.get_mut::<B>().is_none());
+    }
+
     /// Regression test for issue #14348
     #[test]
     fn builder_static_dense_dynamic_sparse() {