Add EntityRefExcept and EntityMutExcept world queries, in preparation for generalized animation. (#15207)

This commit adds two new `WorldQuery` types: `EntityRefExcept` and
`EntityMutExcept`. These types work just like `EntityRef` and
`EntityMut`, but they prevent access to a statically-specified list of
components. For example, `EntityMutExcept<(AnimationPlayer,
Handle<AnimationGraph>)>` provides mutable access to all components
except for `AnimationPlayer` and `Handle<AnimationGraph>`. These types
are useful when you need to be able to process arbitrary queries while
iterating over the results of another `EntityMut` query.

The motivating use case is *generalized animation*, which is an upcoming
feature that allows animation of any component property, not just
rotation, translation, scaling, or morph weights. To implement this, we
must change the current `AnyOf<(&mut Transform, &mut MorphWeights)>` to
instead be `EntityMutExcept<(AnimationPlayer, Handle<AnimationGraph>)>`.
It's possible to use `FilteredEntityMut` in conjunction with a
dynamically-generated system instead, but `FilteredEntityMut` isn't
optimized for the use case of a large number of allowed components
coupled with a small set of disallowed components. No amount of
optimization of `FilteredEntityMut` produced acceptable performance on
the `many_foxes` benchmark. `Query<EntityMut, Without<AnimationPlayer>>`
will not suffice either, as it's legal and idiomatic for an
`AnimationTarget` and an `AnimationPlayer` to coexist on the same
entity.

An alternate proposal was to implement a somewhat-more-general
`Except<Q, CL>` feature, where Q is a `WorldQuery` and CL is a
`ComponentList`. I wasn't able to implement that proposal in a
reasonable way, because of the fact that methods like
`EntityMut::get_mut` and `EntityRef::get` are inherent methods instead
of methods on `WorldQuery`, and therefore there was no way to delegate
methods like `get` and `get_mut` to the inner query in a generic way.
Refactoring those methods into a trait would probably be possible.
However, I didn't see a use case for a hypothetical `Except` with
arbitrary queries: `Query<Except<(&Transform, &Visibility),
Visibility>>` would just be a complicated equivalent to
`Query<&Transform>`, for instance. So, out of a desire for simplicity, I
omitted a generic `Except` mechanism.

I've tested the performance of generalized animation on `many_foxes` and
found that, with this patch, `animate_targets` has a 7.4% slowdown over
`main`. With `FilteredEntityMut` optimized to use `Arc<Access>`, the
slowdown is 75.6%, due to contention on the reference count. Without
`Arc<Access>`, the slowdown is even worse, over 2x.

## Testing

New tests have been added that check that `EntityRefExcept` and
`EntityMutExcept` allow and disallow access to components properly and
that the query engine can correctly reject conflicting queries involving
those types.

A Tracy profile of `many_foxes` with 10,000 foxes showing generalized
animation using `FilteredEntityMut` (red) vs. main (yellow) is as
follows:

![Screenshot 2024-09-12
225914](https://github.com/user-attachments/assets/2993d74c-a513-4ba4-85bd-225672e7170a)

A Tracy profile of `many_foxes` with 10,000 foxes showing generalized
animation using this `EntityMutExcept` (yellow) vs. main (red) is as
follows:

![Screenshot 2024-09-14
205831](https://github.com/user-attachments/assets/4241015e-0c5d-44ef-835b-43f78a24e604)

Author: pcwalton

crates/bevy_ecs/Cargo.toml

@@ -34,6 +34,7 @@ serde = { version = "1", optional = true, default-features = false }
 thiserror = "1.0"
 nonmax = "0.5"
 arrayvec = { version = "0.7.4", optional = true }
+smallvec = "1"
 
 [dev-dependencies]
 rand = "0.8"

crates/bevy_ecs/src/archetype.rs

@@ -499,6 +499,16 @@ impl Archetype {
         self.components.len()
     }
 
+    /// Gets an iterator of all of the components in the archetype, along with
+    /// their archetype component ID.
+    pub(crate) fn components_with_archetype_component_id(
+        &self,
+    ) -> impl Iterator<Item = (ComponentId, ArchetypeComponentId)> + '_ {
+        self.components
+            .iter()
+            .map(|(component_id, info)| (*component_id, info.archetype_component_id))
+    }
+
     /// Fetches an immutable reference to the archetype's [`Edges`], a cache of
     /// archetypal relationships.
     #[inline]

crates/bevy_ecs/src/query/access.rs

@@ -79,10 +79,14 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
                 .map_or(false, |info| info.storage_type() == StorageType::Table)
         };
 
-        self.access
-            .access()
-            .component_reads_and_writes()
-            .all(is_dense)
+        #[allow(deprecated)]
+        let (mut component_reads_and_writes, component_reads_and_writes_inverted) =
+            self.access.access().component_reads_and_writes();
+        if component_reads_and_writes_inverted {
+            return false;
+        }
+
+        component_reads_and_writes.all(is_dense)
             && self.access.access().archetypal().all(is_dense)
             && !self.access.access().has_read_all_components()
             && self.access.with_filters().all(is_dense)

crates/bevy_ecs/src/query/builder.rs

@@ -1,17 +1,19 @@
 use crate::{
     archetype::{Archetype, Archetypes},
+    bundle::Bundle,
     change_detection::{MaybeThinSlicePtrLocation, Ticks, TicksMut},
     component::{Component, ComponentId, Components, StorageType, Tick},
     entity::{Entities, Entity, EntityLocation},
     query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery},
     storage::{ComponentSparseSet, Table, TableRow},
     world::{
-        unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityRef, FilteredEntityMut,
-        FilteredEntityRef, Mut, Ref, World,
+        unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept,
+        FilteredEntityMut, FilteredEntityRef, Mut, Ref, World,
     },
 };
 use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref};
 use bevy_utils::all_tuples;
+use smallvec::SmallVec;
 use std::{cell::UnsafeCell, marker::PhantomData};
 
 /// Types that can be fetched from a [`World`] using a [`Query`].
@@ -626,27 +628,15 @@ unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> {
     unsafe fn set_archetype<'w>(
         fetch: &mut Self::Fetch<'w>,
         state: &Self::State,
-        archetype: &'w Archetype,
+        _: &'w Archetype,
         _table: &Table,
     ) {
-        let mut access = Access::default();
-        state.access.component_reads().for_each(|id| {
-            if archetype.contains(id) {
-                access.add_component_read(id);
-            }
-        });
-        fetch.1 = access;
+        fetch.1.clone_from(&state.access);
     }
 
     #[inline]
-    unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) {
-        let mut access = Access::default();
-        state.access.component_reads().for_each(|id| {
-            if table.has_column(id) {
-                access.add_component_read(id);
-            }
-        });
-        fetch.1 = access;
+    unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, _: &'w Table) {
+        fetch.1.clone_from(&state.access);
     }
 
     #[inline]
@@ -733,37 +723,15 @@ unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> {
     unsafe fn set_archetype<'w>(
         fetch: &mut Self::Fetch<'w>,
         state: &Self::State,
-        archetype: &'w Archetype,
+        _: &'w Archetype,
         _table: &Table,
     ) {
-        let mut access = Access::default();
-        state.access.component_reads().for_each(|id| {
-            if archetype.contains(id) {
-                access.add_component_read(id);
-            }
-        });
-        state.access.component_writes().for_each(|id| {
-            if archetype.contains(id) {
-                access.add_component_write(id);
-            }
-        });
-        fetch.1 = access;
+        fetch.1.clone_from(&state.access);
     }
 
     #[inline]
-    unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) {
-        let mut access = Access::default();
-        state.access.component_reads().for_each(|id| {
-            if table.has_column(id) {
-                access.add_component_read(id);
-            }
-        });
-        state.access.component_writes().for_each(|id| {
-            if table.has_column(id) {
-                access.add_component_write(id);
-            }
-        });
-        fetch.1 = access;
+    unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, _: &'w Table) {
+        fetch.1.clone_from(&state.access);
     }
 
     #[inline]
@@ -815,6 +783,201 @@ unsafe impl<'a> QueryData for FilteredEntityMut<'a> {
     type ReadOnly = FilteredEntityRef<'a>;
 }
 
+/// SAFETY: `EntityRefExcept` guards access to all components in the bundle `B`
+/// and populates `Access` values so that queries that conflict with this access
+/// are rejected.
+unsafe impl<'a, B> WorldQuery for EntityRefExcept<'a, B>
+where
+    B: Bundle,
+{
+    type Fetch<'w> = UnsafeWorldCell<'w>;
+    type Item<'w> = EntityRefExcept<'w, B>;
+    type State = SmallVec<[ComponentId; 4]>;
+
+    fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> {
+        item
+    }
+
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
+    unsafe fn init_fetch<'w>(
+        world: UnsafeWorldCell<'w>,
+        _: &Self::State,
+        _: Tick,
+        _: Tick,
+    ) -> Self::Fetch<'w> {
+        world
+    }
+
+    const IS_DENSE: bool = true;
+
+    unsafe fn set_archetype<'w>(
+        _: &mut Self::Fetch<'w>,
+        _: &Self::State,
+        _: &'w Archetype,
+        _: &'w Table,
+    ) {
+    }
+
+    unsafe fn set_table<'w>(_: &mut Self::Fetch<'w>, _: &Self::State, _: &'w Table) {}
+
+    unsafe fn fetch<'w>(
+        world: &mut Self::Fetch<'w>,
+        entity: Entity,
+        _: TableRow,
+    ) -> Self::Item<'w> {
+        let cell = world.get_entity(entity).unwrap();
+        EntityRefExcept::new(cell)
+    }
+
+    fn update_component_access(
+        state: &Self::State,
+        filtered_access: &mut FilteredAccess<ComponentId>,
+    ) {
+        let mut my_access = Access::new();
+        my_access.read_all_components();
+        for id in state {
+            my_access.remove_component_read(*id);
+        }
+
+        let access = filtered_access.access_mut();
+        assert!(
+            access.is_compatible(&my_access),
+            "`EntityRefExcept<{}>` conflicts with a previous access in this query.",
+            std::any::type_name::<B>(),
+        );
+        access.extend(&my_access);
+    }
+
+    fn init_state(world: &mut World) -> Self::State {
+        Self::get_state(world.components()).unwrap()
+    }
+
+    fn get_state(components: &Components) -> Option<Self::State> {
+        let mut ids = SmallVec::new();
+        B::get_component_ids(components, &mut |maybe_id| {
+            if let Some(id) = maybe_id {
+                ids.push(id);
+            }
+        });
+        Some(ids)
+    }
+
+    fn matches_component_set(_: &Self::State, _: &impl Fn(ComponentId) -> bool) -> bool {
+        true
+    }
+}
+
+/// SAFETY: `Self` is the same as `Self::ReadOnly`.
+unsafe impl<'a, B> QueryData for EntityRefExcept<'a, B>
+where
+    B: Bundle,
+{
+    type ReadOnly = Self;
+}
+
+/// SAFETY: `EntityRefExcept` enforces read-only access to its contained
+/// components.
+unsafe impl<'a, B> ReadOnlyQueryData for EntityRefExcept<'a, B> where B: Bundle {}
+
+/// SAFETY: `EntityMutExcept` guards access to all components in the bundle `B`
+/// and populates `Access` values so that queries that conflict with this access
+/// are rejected.
+unsafe impl<'a, B> WorldQuery for EntityMutExcept<'a, B>
+where
+    B: Bundle,
+{
+    type Fetch<'w> = UnsafeWorldCell<'w>;
+    type Item<'w> = EntityMutExcept<'w, B>;
+    type State = SmallVec<[ComponentId; 4]>;
+
+    fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> {
+        item
+    }
+
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
+    unsafe fn init_fetch<'w>(
+        world: UnsafeWorldCell<'w>,
+        _: &Self::State,
+        _: Tick,
+        _: Tick,
+    ) -> Self::Fetch<'w> {
+        world
+    }
+
+    const IS_DENSE: bool = true;
+
+    unsafe fn set_archetype<'w>(
+        _: &mut Self::Fetch<'w>,
+        _: &Self::State,
+        _: &'w Archetype,
+        _: &'w Table,
+    ) {
+    }
+
+    unsafe fn set_table<'w>(_: &mut Self::Fetch<'w>, _: &Self::State, _: &'w Table) {}
+
+    unsafe fn fetch<'w>(
+        world: &mut Self::Fetch<'w>,
+        entity: Entity,
+        _: TableRow,
+    ) -> Self::Item<'w> {
+        let cell = world.get_entity(entity).unwrap();
+        EntityMutExcept::new(cell)
+    }
+
+    fn update_component_access(
+        state: &Self::State,
+        filtered_access: &mut FilteredAccess<ComponentId>,
+    ) {
+        let mut my_access = Access::new();
+        my_access.write_all_components();
+        for id in state {
+            my_access.remove_component_read(*id);
+        }
+
+        let access = filtered_access.access_mut();
+        assert!(
+            access.is_compatible(&my_access),
+            "`EntityMutExcept<{}>` conflicts with a previous access in this query.",
+            std::any::type_name::<B>()
+        );
+        access.extend(&my_access);
+    }
+
+    fn init_state(world: &mut World) -> Self::State {
+        Self::get_state(world.components()).unwrap()
+    }
+
+    fn get_state(components: &Components) -> Option<Self::State> {
+        let mut ids = SmallVec::new();
+        B::get_component_ids(components, &mut |maybe_id| {
+            if let Some(id) = maybe_id {
+                ids.push(id);
+            }
+        });
+        Some(ids)
+    }
+
+    fn matches_component_set(_: &Self::State, _: &impl Fn(ComponentId) -> bool) -> bool {
+        true
+    }
+}
+
+/// SAFETY: All accesses that `EntityRefExcept` provides are also accesses that
+/// `EntityMutExcept` provides.
+unsafe impl<'a, B> QueryData for EntityMutExcept<'a, B>
+where
+    B: Bundle,
+{
+    type ReadOnly = EntityRefExcept<'a, B>;
+}
+
 /// SAFETY:
 /// `update_component_access` and `update_archetype_component_access` do nothing.
 /// This is sound because `fetch` does not access components.

crates/bevy_ecs/src/query/fetch.rs

@@ -508,22 +508,46 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
         archetype: &Archetype,
         access: &mut Access<ArchetypeComponentId>,
     ) {
-        self.component_access
-            .access
-            .component_reads()
-            .for_each(|id| {
+        // As a fast path, we can iterate directly over the components involved
+        // if the `access` isn't inverted.
+        #[allow(deprecated)]
+        let (component_reads_and_writes, component_reads_and_writes_inverted) =
+            self.component_access.access.component_reads_and_writes();
+        let (component_writes, component_writes_inverted) =
+            self.component_access.access.component_writes();
+
+        if !component_reads_and_writes_inverted && !component_writes_inverted {
+            component_reads_and_writes.for_each(|id| {
                 if let Some(id) = archetype.get_archetype_component_id(id) {
                     access.add_component_read(id);
                 }
             });
-        self.component_access
-            .access
-            .component_writes()
-            .for_each(|id| {
+            component_writes.for_each(|id| {
                 if let Some(id) = archetype.get_archetype_component_id(id) {
                     access.add_component_write(id);
                 }
             });
+            return;
+        }
+
+        for (component_id, archetype_component_id) in
+            archetype.components_with_archetype_component_id()
+        {
+            if self
+                .component_access
+                .access
+                .has_component_read(component_id)
+            {
+                access.add_component_read(archetype_component_id);
+            }
+            if self
+                .component_access
+                .access
+                .has_component_write(component_id)
+            {
+                access.add_component_write(archetype_component_id);
+            }
+        }
     }
 
     /// Use this to transform a [`QueryState`] into a more generic [`QueryState`].