Combine Query and QueryLens using a type parameter for state

crates/bevy_ecs/src/query/fetch.rs

@@ -3,7 +3,7 @@ use crate::{
     bundle::Bundle,
     change_detection::{MaybeLocation, Ticks, TicksMut},
     component::{Component, ComponentId, Components, Mutable, StorageType, Tick},
-    entity::{Entities, Entity, EntityLocation},
+    entity::{Entities, Entity, EntityEquivalent, EntityLocation},
     query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery},
     storage::{ComponentSparseSet, Table, TableRow},
     world::{
@@ -332,6 +332,20 @@ pub unsafe trait QueryData: WorldQuery {
 /// This must only be implemented for read-only [`QueryData`]'s.
 pub unsafe trait ReadOnlyQueryData: QueryData<ReadOnly = Self> {}
 
+/// A [`QueryData`] type that produces a [`EntityEquivalent`] item
+/// equaling the `Entity` it is addressed by.
+///
+/// # Safety
+///
+/// [`<Self as QueryData>::fetch`](QueryData::fetch) must always return an item whose `Entity`
+/// equals the one the function was called with.
+/// I.e.: `Self::fetch(fetch, entity, table_row).entity() == entity` always holds.
+pub unsafe trait EntityEquivalentQueryData: QueryData
+where
+    for<'a> Self: QueryData<Item<'a>: EntityEquivalent>,
+{
+}
+
 /// The item type returned when a [`WorldQuery`] is iterated over
 pub type QueryItem<'w, Q> = <Q as QueryData>::Item<'w>;
 /// The read-only variant of the item type returned when a [`QueryData`] is iterated over immutably
@@ -409,6 +423,9 @@ unsafe impl QueryData for Entity {
 /// SAFETY: access is read only
 unsafe impl ReadOnlyQueryData for Entity {}
 
+/// SAFETY: `entity()` returns `self`, and `fetch` returns `entity`
+unsafe impl EntityEquivalentQueryData for Entity {}
+
 /// SAFETY:
 /// `update_component_access` and `update_archetype_component_access` do nothing.
 /// This is sound because `fetch` does not access components.
@@ -570,6 +587,9 @@ unsafe impl<'a> QueryData for EntityRef<'a> {
 /// SAFETY: access is read only
 unsafe impl ReadOnlyQueryData for EntityRef<'_> {}
 
+/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity`
+unsafe impl EntityEquivalentQueryData for EntityRef<'_> {}
+
 /// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self`
 unsafe impl<'a> WorldQuery for EntityMut<'a> {
     type Fetch<'w> = UnsafeWorldCell<'w>;
@@ -648,6 +668,9 @@ unsafe impl<'a> QueryData for EntityMut<'a> {
     }
 }
 
+/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity`
+unsafe impl EntityEquivalentQueryData for EntityMut<'_> {}
+
 /// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self`
 unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> {
     type Fetch<'w> = (UnsafeWorldCell<'w>, Access<ComponentId>);
@@ -757,6 +780,9 @@ unsafe impl<'a> QueryData for FilteredEntityRef<'a> {
 /// SAFETY: Access is read-only.
 unsafe impl ReadOnlyQueryData for FilteredEntityRef<'_> {}
 
+/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity`
+unsafe impl EntityEquivalentQueryData for FilteredEntityRef<'_> {}
+
 /// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self`
 unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> {
     type Fetch<'w> = (UnsafeWorldCell<'w>, Access<ComponentId>);
@@ -861,6 +887,9 @@ unsafe impl<'a> QueryData for FilteredEntityMut<'a> {
     }
 }
 
+/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity`
+unsafe impl EntityEquivalentQueryData for FilteredEntityMut<'_> {}
+
 /// 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.
@@ -961,6 +990,9 @@ where
 /// components.
 unsafe impl<'a, B> ReadOnlyQueryData for EntityRefExcept<'a, B> where B: Bundle {}
 
+/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity`
+unsafe impl<B: Bundle> EntityEquivalentQueryData for EntityRefExcept<'_, B> {}
+
 /// 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.
@@ -1058,6 +1090,9 @@ where
     }
 }
 
+/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity`
+unsafe impl<B: Bundle> EntityEquivalentQueryData for EntityMutExcept<'_, B> {}
+
 /// SAFETY:
 /// `update_component_access` and `update_archetype_component_access` do nothing.
 /// This is sound because `fetch` does not access components.