Implement QueryData for &Archetype and EntityLocation (#12398)

# Objective
Fixes #12392, fixes #12393, and fixes #11387. Implement QueryData for
Archetype and EntityLocation.

## Solution
Add impls for both of the types.

---

## Changelog
Added: `&Archetype` now implements `QueryData`
Added: `EntityLocation` now implements `QueryData`

---------

Co-authored-by: Alice Cecile <[email protected]>

crates/bevy_ecs/src/query/fetch.rs

@@ -1,8 +1,8 @@
 use crate::{
-    archetype::Archetype,
+    archetype::{Archetype, Archetypes},
     change_detection::{Ticks, TicksMut},
     component::{Component, ComponentId, StorageType, Tick},
-    entity::Entity,
+    entity::{Entities, Entity, EntityLocation},
     query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery},
     storage::{ComponentSparseSet, Table, TableRow},
     world::{
@@ -18,7 +18,7 @@ use std::{cell::UnsafeCell, marker::PhantomData};
 ///
 /// There are many types that natively implement this trait:
 ///
-/// - **Component references.**
+/// - **Component references. (&T and &mut T)**
 ///   Fetches a component by reference (immutably or mutably).
 /// - **`QueryData` tuples.**
 ///   If every element of a tuple implements `QueryData`, then the tuple itself also implements the same trait.
@@ -27,6 +27,14 @@ use std::{cell::UnsafeCell, marker::PhantomData};
 ///   but nesting of tuples allows infinite `WorldQuery`s.
 /// - **[`Entity`].**
 ///   Gets the identifier of the queried entity.
+/// - **[`EntityLocation`].**
+///   Gets the location metadata of the queried entity.
+/// - **[`EntityRef`].**
+///   Read-only access to arbitrary components on the queried entity.
+/// - **[`EntityMut`].**
+///   Mutable access to arbitrary components on the queried entity.
+/// - **[`&Archetype`](Archetype).**
+///   Read-only access to the archetype-level metadata of the queried entity.
 /// - **[`Option`].**
 ///   By default, a world query only tests entities that have the matching component types.
 ///   Wrapping it into an `Option` will increase the query search space, and it will return `None` if an entity doesn't satisfy the `WorldQuery`.
@@ -345,6 +353,78 @@ unsafe impl QueryData for Entity {
 /// SAFETY: access is read only
 unsafe impl ReadOnlyQueryData for Entity {}
 
+/// SAFETY:
+/// `update_component_access` and `update_archetype_component_access` do nothing.
+/// This is sound because `fetch` does not access components.
+unsafe impl WorldQuery for EntityLocation {
+    type Item<'w> = EntityLocation;
+    type Fetch<'w> = &'w Entities;
+    type State = ();
+
+    fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> {
+        item
+    }
+
+    unsafe fn init_fetch<'w>(
+        world: UnsafeWorldCell<'w>,
+        _state: &Self::State,
+        _last_run: Tick,
+        _this_run: Tick,
+    ) -> Self::Fetch<'w> {
+        world.entities()
+    }
+
+    // This is set to true to avoid forcing archetypal iteration in compound queries, is likely to be slower
+    // in most practical use case.
+    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> {
+        // SAFETY: `fetch` must be called with an entity that exists in the world
+        unsafe { fetch.get(entity).debug_checked_unwrap() }
+    }
+
+    fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess<ComponentId>) {}
+
+    fn init_state(_world: &mut World) {}
+
+    fn get_state(_world: &World) -> Option<()> {
+        Some(())
+    }
+
+    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 EntityLocation {
+    type ReadOnly = Self;
+}
+
+/// SAFETY: access is read only
+unsafe impl ReadOnlyQueryData for EntityLocation {}
+
 /// SAFETY:
 /// `fetch` accesses all components in a readonly way.
 /// This is sound because `update_component_access` and `update_archetype_component_access` set read access for all components and panic when appropriate.
@@ -709,6 +789,81 @@ unsafe impl<'a> QueryData for FilteredEntityMut<'a> {
     type ReadOnly = FilteredEntityRef<'a>;
 }
 
+/// SAFETY:
+/// `update_component_access` and `update_archetype_component_access` do nothing.
+/// This is sound because `fetch` does not access components.
+unsafe impl WorldQuery for &Archetype {
+    type Item<'w> = &'w Archetype;
+    type Fetch<'w> = (&'w Entities, &'w Archetypes);
+    type State = ();
+
+    fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> {
+        item
+    }
+
+    unsafe fn init_fetch<'w>(
+        world: UnsafeWorldCell<'w>,
+        _state: &Self::State,
+        _last_run: Tick,
+        _this_run: Tick,
+    ) -> Self::Fetch<'w> {
+        (world.entities(), world.archetypes())
+    }
+
+    // This could probably be a non-dense query and just set a Option<&Archetype> fetch value in
+    // set_archetypes, but forcing archetypal iteration is likely to be slower in any compound query.
+    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> {
+        let (entities, archetypes) = *fetch;
+        // SAFETY: `fetch` must be called with an entity that exists in the world
+        let location = unsafe { entities.get(entity).debug_checked_unwrap() };
+        // SAFETY: The assigned archetype for a living entity must always be valid.
+        unsafe { archetypes.get(location.archetype_id).debug_checked_unwrap() }
+    }
+
+    fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess<ComponentId>) {}
+
+    fn init_state(_world: &mut World) {}
+
+    fn get_state(_world: &World) -> Option<()> {
+        Some(())
+    }
+
+    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 &Archetype {
+    type ReadOnly = Self;
+}
+
+/// SAFETY: access is read only
+unsafe impl ReadOnlyQueryData for &Archetype {}
+
 #[doc(hidden)]
 pub struct ReadFetch<'w, T> {
     // T::STORAGE_TYPE = StorageType::Table

crates/bevy_ecs/src/system/query.rs

@@ -1333,12 +1333,16 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     /// Besides removing parameters from the query, you can also
     /// make limited changes to the types of parameters.
     ///
-    /// * Can always add/remove `Entity`
+    /// * Can always add/remove [`Entity`]
+    /// * Can always add/remove [`EntityLocation`]
+    /// * Can always add/remove [`&Archetype`]
     /// * `Ref<T>` <-> `&T`
     /// * `&mut T` -> `&T`
     /// * `&mut T` -> `Ref<T>`
     /// * [`EntityMut`](crate::world::EntityMut) -> [`EntityRef`](crate::world::EntityRef)
-    ///    
+    ///  
+    /// [`EntityLocation`]: crate::entity::EntityLocation
+    /// [`&Archetype`]: crate::archetype::Archetype
     pub fn transmute_lens<NewD: QueryData>(&mut self) -> QueryLens<'_, NewD> {
         self.transmute_lens_filtered::<NewD, ()>()
     }