Combine Query and QueryLens using a type parameter for state, but don't introduce owned iterators

crates/bevy_ecs/src/query/state.rs

@@ -14,9 +14,9 @@ use crate::{
 #[cfg(all(not(target_arch = "wasm32"), feature = "multi_threaded"))]
 use crate::entity::UniqueEntityEquivalentSlice;
 
-use alloc::vec::Vec;
+use alloc::{boxed::Box, vec::Vec};
 use bevy_utils::prelude::DebugName;
-use core::{fmt, ptr};
+use core::{fmt, ops::Deref, ptr};
 use fixedbitset::FixedBitSet;
 use log::warn;
 #[cfg(feature = "trace")]
@@ -89,6 +89,30 @@ pub struct QueryState<D: QueryData, F: QueryFilter = ()> {
     par_iter_span: Span,
 }
 
+/// Indicates how the internal state of a [`Query`] is stored.
+pub trait QueryType {
+    /// The type used to store the internal state of a [`Query`].
+    type QueryState<'s, D: QueryData + 's, F: QueryFilter + 's>: Deref<Target = QueryState<D, F>>;
+}
+
+/// A [`Query`] that holds a shared borrow of its internal state.
+///
+/// This is the default [`QueryType`] for a [`Query`], and is used by [system parameters](crate::system::SystemParam).
+pub struct BorrowedQuery;
+
+impl QueryType for BorrowedQuery {
+    type QueryState<'s, D: QueryData + 's, F: QueryFilter + 's> = &'s QueryState<D, F>;
+}
+
+/// A [`Query`] that owns its internal state.
+///
+/// This is the [`QueryType`] for a [`QueryLens`](crate::system::QueryLens).
+pub struct OwnedQuery;
+
+impl QueryType for OwnedQuery {
+    type QueryState<'s, D: QueryData + 's, F: QueryFilter + 's> = Box<QueryState<D, F>>;
+}
+
 impl<D: QueryData, F: QueryFilter> fmt::Debug for QueryState<D, F> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("QueryState")
@@ -2043,7 +2067,7 @@ mod tests {
             .run_system_once(|query: Query<&mut A>| {
                 let mut readonly = query.as_readonly();
                 let mut lens: QueryLens<&mut A> = readonly.transmute_lens();
-                bad(lens.query(), query.as_readonly());
+                bad(lens.reborrow(), query.as_readonly());
             })
             .unwrap();
     }
@@ -2222,7 +2246,7 @@ mod tests {
             .run_system_once(|query_a: Query<&mut A>, mut query_b: Query<&mut B>| {
                 let mut readonly = query_a.as_readonly();
                 let mut lens: QueryLens<(&mut A, &mut B)> = readonly.join(&mut query_b);
-                bad(lens.query(), query_a.as_readonly());
+                bad(lens.reborrow(), query_a.as_readonly());
             })
             .unwrap();
     }

crates/bevy_ecs/src/relationship/relationship_query.rs

@@ -1,6 +1,6 @@
 use crate::{
     entity::Entity,
-    query::{QueryData, QueryFilter},
+    query::{QueryData, QueryFilter, QueryType},
     relationship::{Relationship, RelationshipTarget},
     system::Query,
 };
@@ -9,12 +9,12 @@ use smallvec::SmallVec;
 
 use super::SourceIter;
 
-impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType> Query<'w, 's, D, F, T> {
     /// If the given `entity` contains the `R` [`Relationship`] component, returns the
     /// target entity of that relationship.
     pub fn related<R: Relationship>(&'w self, entity: Entity) -> Option<Entity>
     where
-        <D as QueryData>::ReadOnly: QueryData<Item<'w, 's> = &'w R>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w R>,
     {
         self.get(entity).map(R::get).ok()
     }
@@ -26,7 +26,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
         entity: Entity,
     ) -> impl Iterator<Item = Entity> + 'w
     where
-        <D as QueryData>::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
     {
         self.get(entity)
             .into_iter()
@@ -42,7 +42,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity.
     pub fn root_ancestor<R: Relationship>(&'w self, entity: Entity) -> Entity
     where
-        <D as QueryData>::ReadOnly: QueryData<Item<'w, 's> = &'w R>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w R>,
     {
         // Recursively search up the tree until we're out of parents
         match self.get(entity) {
@@ -60,9 +60,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     pub fn iter_leaves<S: RelationshipTarget>(
         &'w self,
         entity: Entity,
-    ) -> impl Iterator<Item = Entity> + use<'w, 's, S, D, F>
+    ) -> impl Iterator<Item = Entity> + use<'w, 's, S, D, F, T>
     where
-        <D as QueryData>::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
         SourceIter<'w, S>: DoubleEndedIterator,
     {
         self.iter_descendants_depth_first(entity).filter(|entity| {
@@ -80,7 +80,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
         entity: Entity,
     ) -> impl Iterator<Item = Entity> + 'w
     where
-        D::ReadOnly: QueryData<Item<'w, 's> = (Option<&'w R>, Option<&'w R::RelationshipTarget>)>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = (Option<&'w R>, Option<&'w R::RelationshipTarget>)>,
     {
         self.get(entity)
             .ok()
@@ -101,9 +101,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     pub fn iter_descendants<S: RelationshipTarget>(
         &'w self,
         entity: Entity,
-    ) -> DescendantIter<'w, 's, D, F, S>
+    ) -> DescendantIter<'w, 's, D, F, T, S>
     where
-        D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
     {
         DescendantIter::new(self, entity)
     }
@@ -118,9 +118,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     pub fn iter_descendants_depth_first<S: RelationshipTarget>(
         &'w self,
         entity: Entity,
-    ) -> DescendantDepthFirstIter<'w, 's, D, F, S>
+    ) -> DescendantDepthFirstIter<'w, 's, D, F, T, S>
     where
-        D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
         SourceIter<'w, S>: DoubleEndedIterator,
     {
         DescendantDepthFirstIter::new(self, entity)
@@ -135,9 +135,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     pub fn iter_ancestors<R: Relationship>(
         &'w self,
         entity: Entity,
-    ) -> AncestorIter<'w, 's, D, F, R>
+    ) -> AncestorIter<'w, 's, D, F, T, R>
     where
-        D::ReadOnly: QueryData<Item<'w, 's> = &'w R>,
+        D::ReadOnly: QueryData<Item<'w, 'w> = &'w R>,
     {
         AncestorIter::new(self, entity)
     }
@@ -146,20 +146,21 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
 /// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
 ///
 /// Traverses the hierarchy breadth-first.
-pub struct DescendantIter<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
+pub struct DescendantIter<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, S: RelationshipTarget>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
 {
-    children_query: &'w Query<'w, 's, D, F>,
+    children_query: &'w Query<'w, 's, D, F, T>,
     vecdeque: VecDeque<Entity>,
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> DescendantIter<'w, 's, D, F, S>
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, S: RelationshipTarget>
+    DescendantIter<'w, 's, D, F, T, S>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
 {
     /// Returns a new [`DescendantIter`].
-    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
+    pub fn new(children_query: &'w Query<'w, 's, D, F, T>, entity: Entity) -> Self {
         DescendantIter {
             children_query,
             vecdeque: children_query
@@ -171,10 +172,10 @@ where
     }
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator
-    for DescendantIter<'w, 's, D, F, S>
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, S: RelationshipTarget> Iterator
+    for DescendantIter<'w, 's, D, F, T, S>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
 {
     type Item = Entity;
 
@@ -192,22 +193,28 @@ where
 /// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
 ///
 /// Traverses the hierarchy depth-first.
-pub struct DescendantDepthFirstIter<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
-where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+pub struct DescendantDepthFirstIter<
+    'w,
+    's,
+    D: QueryData,
+    F: QueryFilter,
+    T: QueryType,
+    S: RelationshipTarget,
+> where
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
 {
-    children_query: &'w Query<'w, 's, D, F>,
+    children_query: &'w Query<'w, 's, D, F, T>,
     stack: SmallVec<[Entity; 8]>,
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
-    DescendantDepthFirstIter<'w, 's, D, F, S>
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, S: RelationshipTarget>
+    DescendantDepthFirstIter<'w, 's, D, F, T, S>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
     SourceIter<'w, S>: DoubleEndedIterator,
 {
     /// Returns a new [`DescendantDepthFirstIter`].
-    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
+    pub fn new(children_query: &'w Query<'w, 's, D, F, T>, entity: Entity) -> Self {
         DescendantDepthFirstIter {
             children_query,
             stack: children_query
@@ -217,10 +224,10 @@ where
     }
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator
-    for DescendantDepthFirstIter<'w, 's, D, F, S>
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, S: RelationshipTarget> Iterator
+    for DescendantDepthFirstIter<'w, 's, D, F, T, S>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w S>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w S>,
     SourceIter<'w, S>: DoubleEndedIterator,
 {
     type Item = Entity;
@@ -237,31 +244,32 @@ where
 }
 
 /// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`].
-pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter, R: Relationship>
+pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, R: Relationship>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w R>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w R>,
 {
-    parent_query: &'w Query<'w, 's, D, F>,
+    parent_query: &'w Query<'w, 's, D, F, T>,
     next: Option<Entity>,
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> AncestorIter<'w, 's, D, F, R>
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, R: Relationship>
+    AncestorIter<'w, 's, D, F, T, R>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w R>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w R>,
 {
     /// Returns a new [`AncestorIter`].
-    pub fn new(parent_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
+    pub fn new(parent_query: &'w Query<'w, 's, D, F, T>, entity: Entity) -> Self {
         AncestorIter {
             parent_query,
             next: Some(entity),
         }
     }
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> Iterator
-    for AncestorIter<'w, 's, D, F, R>
+impl<'w, 's, D: QueryData, F: QueryFilter, T: QueryType, R: Relationship> Iterator
+    for AncestorIter<'w, 's, D, F, T, R>
 where
-    D::ReadOnly: QueryData<Item<'w, 's> = &'w R>,
+    D::ReadOnly: QueryData<Item<'w, 'w> = &'w R>,
 {
     type Item = Entity;