Add 's (state) lifetime to Fetch (#2515)

Allows iterators to return things that borrow data from `QueryState`, needed this in my relations PR figure might be worth landing separately maybe

crates/bevy_ecs/src/query/fetch.rs

@@ -41,11 +41,11 @@ use std::{
 ///
 /// [`Or`]: crate::query::Or
 pub trait WorldQuery {
-    type Fetch: for<'a> Fetch<'a, State = Self::State>;
+    type Fetch: for<'world, 'state> Fetch<'world, 'state, State = Self::State>;
     type State: FetchState;
 }
 
-pub trait Fetch<'w>: Sized {
+pub trait Fetch<'world, 'state>: Sized {
     type Item;
     type State: FetchState;
 
@@ -173,7 +173,7 @@ unsafe impl FetchState for EntityState {
     }
 }
 
-impl<'w> Fetch<'w> for EntityFetch {
+impl<'w, 's> Fetch<'w, 's> for EntityFetch {
     type Item = Entity;
     type State = EntityState;
 
@@ -296,7 +296,7 @@ impl<T> Clone for ReadFetch<T> {
 /// SAFETY: access is read only
 unsafe impl<T> ReadOnlyFetch for ReadFetch<T> {}
 
-impl<'w, T: Component> Fetch<'w> for ReadFetch<T> {
+impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
     type Item = &'w T;
     type State = ReadState<T>;
 
@@ -459,7 +459,7 @@ unsafe impl<T: Component> FetchState for WriteState<T> {
     }
 }
 
-impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
+impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
     type Item = Mut<'w, T>;
     type State = WriteState<T>;
 
@@ -619,7 +619,7 @@ unsafe impl<T: FetchState> FetchState for OptionState<T> {
     }
 }
 
-impl<'w, T: Fetch<'w>> Fetch<'w> for OptionFetch<T> {
+impl<'w, 's, T: Fetch<'w, 's>> Fetch<'w, 's> for OptionFetch<T> {
     type Item = Option<T::Item>;
     type State = OptionState<T::State>;
 
@@ -810,7 +810,7 @@ pub struct ChangeTrackersFetch<T> {
 /// SAFETY: access is read only
 unsafe impl<T> ReadOnlyFetch for ChangeTrackersFetch<T> {}
 
-impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
+impl<'w, 's, T: Component> Fetch<'w, 's> for ChangeTrackersFetch<T> {
     type Item = ChangeTrackers<T>;
     type State = ChangeTrackersState<T>;
 
@@ -913,7 +913,7 @@ impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
 macro_rules! impl_tuple_fetch {
     ($(($name: ident, $state: ident)),*) => {
         #[allow(non_snake_case)]
-        impl<'a, $($name: Fetch<'a>),*> Fetch<'a> for ($($name,)*) {
+        impl<'w, 's, $($name: Fetch<'w, 's>),*> Fetch<'w, 's> for ($($name,)*) {
             type Item = ($($name::Item,)*);
             type State = ($($name::State,)*);
 

crates/bevy_ecs/src/query/filter.rs

@@ -12,7 +12,7 @@ use std::{cell::UnsafeCell, marker::PhantomData, ptr};
 
 /// Extension trait for [`Fetch`] containing methods used by query filters.
 /// This trait exists to allow "short circuit" behaviors for relevant query filter fetches.
-pub trait FilterFetch: for<'a> Fetch<'a> {
+pub trait FilterFetch: for<'w, 's> Fetch<'w, 's> {
     /// # Safety
     ///
     /// Must always be called _after_ [`Fetch::set_archetype`]. `archetype_index` must be in the range
@@ -28,7 +28,7 @@ pub trait FilterFetch: for<'a> Fetch<'a> {
 
 impl<T> FilterFetch for T
 where
-    T: for<'a> Fetch<'a, Item = bool>,
+    T: for<'w, 's> Fetch<'w, 's, Item = bool>,
 {
     #[inline]
     unsafe fn archetype_filter_fetch(&mut self, archetype_index: usize) -> bool {
@@ -119,7 +119,7 @@ unsafe impl<T: Component> FetchState for WithState<T> {
     }
 }
 
-impl<'a, T: Component> Fetch<'a> for WithFetch<T> {
+impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch<T> {
     type Item = bool;
     type State = WithState<T>;
 
@@ -238,7 +238,7 @@ unsafe impl<T: Component> FetchState for WithoutState<T> {
     }
 }
 
-impl<'a, T: Component> Fetch<'a> for WithoutFetch<T> {
+impl<'w, 's, T: Component> Fetch<'w, 's> for WithoutFetch<T> {
     type Item = bool;
     type State = WithoutState<T>;
 
@@ -338,7 +338,7 @@ unsafe impl<T: Bundle> FetchState for WithBundleState<T> {
     }
 }
 
-impl<'a, T: Bundle> Fetch<'a> for WithBundleFetch<T> {
+impl<'w, 's, T: Bundle> Fetch<'w, 's> for WithBundleFetch<T> {
     type Item = bool;
     type State = WithBundleState<T>;
 
@@ -446,8 +446,8 @@ macro_rules! impl_query_filter_tuple {
 
         #[allow(unused_variables)]
         #[allow(non_snake_case)]
-        impl<'a, $($filter: FilterFetch),*> Fetch<'a> for Or<($(OrFetch<$filter>,)*)> {
+        impl<'w, 's, $($filter: FilterFetch),*> Fetch<'w, 's> for Or<($(OrFetch<$filter>,)*)> {
-            type State = Or<($(<$filter as Fetch<'a>>::State,)*)>;
+            type State = Or<($(<$filter as Fetch<'w, 's>>::State,)*)>;
             type Item = bool;
 
             unsafe fn init(world: &World, state: &Self::State, last_change_tick: u32, change_tick: u32) -> Self {
@@ -612,7 +612,7 @@ macro_rules! impl_tick_filter {
             }
         }
 
-        impl<'a, T: Component> Fetch<'a> for $fetch_name<T> {
+        impl<'w, 's, T: Component> Fetch<'w, 's> for $fetch_name<T> {
             type State = $state_name<T>;
             type Item = bool;
 

crates/bevy_ecs/src/query/iter.rs

@@ -131,7 +131,7 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Iterator for QueryIter<'w, 's, Q, F>
 where
     F::Fetch: FilterFetch,
 {
-    type Item = <Q::Fetch as Fetch<'w>>::Item;
+    type Item = <Q::Fetch as Fetch<'w, 's>>::Item;
 
     // NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
     // QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
@@ -279,7 +279,7 @@ where
     /// It is always safe for shared access.
     unsafe fn fetch_next_aliased_unchecked<'a>(
         &mut self,
-    ) -> Option<[<Q::Fetch as Fetch<'a>>::Item; K]>
+    ) -> Option<[<Q::Fetch as Fetch<'a, 's>>::Item; K]>
     where
         Q::Fetch: Clone,
         F::Fetch: Clone,
@@ -309,7 +309,7 @@ where
         }
 
         // TODO: use MaybeUninit::uninit_array if it stabilizes
-        let mut values: [MaybeUninit<<Q::Fetch as Fetch<'a>>::Item>; K] =
+        let mut values: [MaybeUninit<<Q::Fetch as Fetch<'a, 's>>::Item>; K] =
             MaybeUninit::uninit().assume_init();
 
         for (value, cursor) in values.iter_mut().zip(&mut self.cursors) {
@@ -317,15 +317,15 @@ where
         }
 
         // TODO: use MaybeUninit::array_assume_init if it stabilizes
-        let values: [<Q::Fetch as Fetch<'a>>::Item; K] =
+        let values: [<Q::Fetch as Fetch<'a, 's>>::Item; K] =
-            (&values as *const _ as *const [<Q::Fetch as Fetch<'a>>::Item; K]).read();
+            (&values as *const _ as *const [<Q::Fetch as Fetch<'a, 's>>::Item; K]).read();
 
         Some(values)
     }
 
     /// Get next combination of queried components
     #[inline]
-    pub fn fetch_next(&mut self) -> Option<[<Q::Fetch as Fetch<'_>>::Item; K]>
+    pub fn fetch_next(&mut self) -> Option<[<Q::Fetch as Fetch<'_, 's>>::Item; K]>
     where
         Q::Fetch: Clone,
         F::Fetch: Clone,
@@ -346,7 +346,7 @@ where
     Q::Fetch: Clone + ReadOnlyFetch,
     F::Fetch: Clone + FilterFetch + ReadOnlyFetch,
 {
-    type Item = [<Q::Fetch as Fetch<'w>>::Item; K];
+    type Item = [<Q::Fetch as Fetch<'w, 's>>::Item; K];
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
@@ -476,7 +476,7 @@ where
 
     /// retrieve item returned from most recent `next` call again.
     #[inline]
-    unsafe fn peek_last<'w>(&mut self) -> Option<<Q::Fetch as Fetch<'w>>::Item> {
+    unsafe fn peek_last<'w>(&mut self) -> Option<<Q::Fetch as Fetch<'w, 's>>::Item> {
         if self.current_index > 0 {
             if self.is_dense {
                 Some(self.fetch.table_fetch(self.current_index - 1))
@@ -497,7 +497,7 @@ where
         tables: &'w Tables,
         archetypes: &'w Archetypes,
         query_state: &'s QueryState<Q, F>,
-    ) -> Option<<Q::Fetch as Fetch<'w>>::Item> {
+    ) -> Option<<Q::Fetch as Fetch<'w, 's>>::Item> {
         if self.is_dense {
             loop {
                 if self.current_index == self.current_len {

crates/bevy_ecs/src/query/state.rs

@@ -121,7 +121,7 @@ where
         &mut self,
         world: &'w World,
         entity: Entity,
-    ) -> Result<<Q::Fetch as Fetch<'w>>::Item, QueryEntityError>
+    ) -> Result<<Q::Fetch as Fetch<'w, '_>>::Item, QueryEntityError>
     where
         Q::Fetch: ReadOnlyFetch,
     {
@@ -134,7 +134,7 @@ where
         &mut self,
         world: &'w mut World,
         entity: Entity,
-    ) -> Result<<Q::Fetch as Fetch<'w>>::Item, QueryEntityError> {
+    ) -> Result<<Q::Fetch as Fetch<'w, '_>>::Item, QueryEntityError> {
         // SAFETY: query has unique world access
         unsafe { self.get_unchecked(world, entity) }
     }
@@ -148,7 +148,7 @@ where
         &mut self,
         world: &'w World,
         entity: Entity,
-    ) -> Result<<Q::Fetch as Fetch<'w>>::Item, QueryEntityError> {
+    ) -> Result<<Q::Fetch as Fetch<'w, '_>>::Item, QueryEntityError> {
         self.validate_world_and_update_archetypes(world);
         self.get_unchecked_manual(
             world,
@@ -167,7 +167,7 @@ where
         entity: Entity,
         last_change_tick: u32,
         change_tick: u32,
-    ) -> Result<<Q::Fetch as Fetch<'w>>::Item, QueryEntityError> {
+    ) -> Result<<Q::Fetch as Fetch<'w, '_>>::Item, QueryEntityError> {
         let location = world
             .entities
             .get(entity)
@@ -290,10 +290,10 @@ where
     }
 
     #[inline]
-    pub fn for_each<'w>(
+    pub fn for_each<'w, 's>(
-        &mut self,
+        &'s mut self,
         world: &'w World,
-        func: impl FnMut(<Q::Fetch as Fetch<'w>>::Item),
+        func: impl FnMut(<Q::Fetch as Fetch<'w, 's>>::Item),
     ) where
         Q::Fetch: ReadOnlyFetch,
     {
@@ -304,10 +304,10 @@ where
     }
 
     #[inline]
-    pub fn for_each_mut<'w>(
+    pub fn for_each_mut<'w, 's>(
-        &mut self,
+        &'s mut self,
         world: &'w mut World,
-        func: impl FnMut(<Q::Fetch as Fetch<'w>>::Item),
+        func: impl FnMut(<Q::Fetch as Fetch<'w, 's>>::Item),
     ) {
         // SAFETY: query has unique world access
         unsafe {
@@ -320,10 +320,10 @@ where
     /// This does not check for mutable query correctness. To be safe, make sure mutable queries
     /// have unique access to the components they query.
     #[inline]
-    pub unsafe fn for_each_unchecked<'w>(
+    pub unsafe fn for_each_unchecked<'w, 's>(
-        &mut self,
+        &'s mut self,
         world: &'w World,
-        func: impl FnMut(<Q::Fetch as Fetch<'w>>::Item),
+        func: impl FnMut(<Q::Fetch as Fetch<'w, 's>>::Item),
     ) {
         self.validate_world_and_update_archetypes(world);
         self.for_each_unchecked_manual(
@@ -335,12 +335,12 @@ where
     }
 
     #[inline]
-    pub fn par_for_each<'w>(
+    pub fn par_for_each<'w, 's>(
-        &mut self,
+        &'s mut self,
         world: &'w World,
         task_pool: &TaskPool,
         batch_size: usize,
-        func: impl Fn(<Q::Fetch as Fetch<'w>>::Item) + Send + Sync + Clone,
+        func: impl Fn(<Q::Fetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone,
     ) where
         Q::Fetch: ReadOnlyFetch,
     {
@@ -351,12 +351,12 @@ where
     }
 
     #[inline]
-    pub fn par_for_each_mut<'w>(
+    pub fn par_for_each_mut<'w, 's>(
-        &mut self,
+        &'s mut self,
         world: &'w mut World,
         task_pool: &TaskPool,
         batch_size: usize,
-        func: impl Fn(<Q::Fetch as Fetch<'w>>::Item) + Send + Sync + Clone,
+        func: impl Fn(<Q::Fetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone,
     ) {
         // SAFETY: query has unique world access
         unsafe {
@@ -369,12 +369,12 @@ where
     /// This does not check for mutable query correctness. To be safe, make sure mutable queries
     /// have unique access to the components they query.
     #[inline]
-    pub unsafe fn par_for_each_unchecked<'w>(
+    pub unsafe fn par_for_each_unchecked<'w, 's>(
-        &mut self,
+        &'s mut self,
         world: &'w World,
         task_pool: &TaskPool,
         batch_size: usize,
-        func: impl Fn(<Q::Fetch as Fetch<'w>>::Item) + Send + Sync + Clone,
+        func: impl Fn(<Q::Fetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone,
     ) {
         self.validate_world_and_update_archetypes(world);
         self.par_for_each_unchecked_manual(
@@ -396,7 +396,7 @@ where
     pub(crate) unsafe fn for_each_unchecked_manual<'w, 's>(
         &'s self,
         world: &'w World,
-        mut func: impl FnMut(<Q::Fetch as Fetch<'w>>::Item),
+        mut func: impl FnMut(<Q::Fetch as Fetch<'w, 's>>::Item),
         last_change_tick: u32,
         change_tick: u32,
     ) {
@@ -450,7 +450,7 @@ where
         world: &'w World,
         task_pool: &TaskPool,
         batch_size: usize,
-        func: impl Fn(<Q::Fetch as Fetch<'w>>::Item) + Send + Sync + Clone,
+        func: impl Fn(<Q::Fetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone,
         last_change_tick: u32,
         change_tick: u32,
     ) {

crates/bevy_ecs/src/system/query.rs

@@ -266,7 +266,7 @@ where
     ///
     /// This can only be called for read-only queries, see [`Self::for_each_mut`] for write-queries.
     #[inline]
-    pub fn for_each(&self, f: impl FnMut(<Q::Fetch as Fetch<'w>>::Item))
+    pub fn for_each<'s>(&'s self, f: impl FnMut(<Q::Fetch as Fetch<'w, 's>>::Item))
     where
         Q::Fetch: ReadOnlyFetch,
     {
@@ -285,7 +285,7 @@ where
     /// Runs `f` on each query result. This is faster than the equivalent iter() method, but cannot
     /// be chained like a normal [`Iterator`].
     #[inline]
-    pub fn for_each_mut(&mut self, f: impl FnMut(<Q::Fetch as Fetch<'w>>::Item)) {
+    pub fn for_each_mut<'s>(&'s mut self, f: impl FnMut(<Q::Fetch as Fetch<'w, 's>>::Item)) {
         // SAFE: system runs without conflicts with other systems. same-system queries have runtime
         // borrow checks when they conflict
         unsafe {
@@ -303,11 +303,11 @@ where
     /// This can only be called for read-only queries, see [`Self::par_for_each_mut`] for
     /// write-queries.
     #[inline]
-    pub fn par_for_each(
+    pub fn par_for_each<'s>(
-        &self,
+        &'s self,
         task_pool: &TaskPool,
         batch_size: usize,
-        f: impl Fn(<Q::Fetch as Fetch<'w>>::Item) + Send + Sync + Clone,
+        f: impl Fn(<Q::Fetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone,
     ) where
         Q::Fetch: ReadOnlyFetch,
     {
@@ -327,11 +327,11 @@ where
 
     /// Runs `f` on each query result in parallel using the given task pool.
     #[inline]
-    pub fn par_for_each_mut(
+    pub fn par_for_each_mut<'s>(
-        &mut self,
+        &'s mut self,
         task_pool: &TaskPool,
         batch_size: usize,
-        f: impl Fn(<Q::Fetch as Fetch<'w>>::Item) + Send + Sync + Clone,
+        f: impl Fn(<Q::Fetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone,
     ) {
         // SAFE: system runs without conflicts with other systems. same-system queries have runtime
         // borrow checks when they conflict
@@ -511,7 +511,7 @@ where
     /// ```
     ///
     /// This can only be called for read-only queries, see [`Self::single_mut`] for write-queries.
-    pub fn single(&self) -> Result<<Q::Fetch as Fetch<'_>>::Item, QuerySingleError>
+    pub fn single(&self) -> Result<<Q::Fetch as Fetch<'_, '_>>::Item, QuerySingleError>
     where
         Q::Fetch: ReadOnlyFetch,
     {
@@ -530,7 +530,7 @@ where
 
     /// Gets the query result if it is only a single result, otherwise returns a
     /// [`QuerySingleError`].
-    pub fn single_mut(&mut self) -> Result<<Q::Fetch as Fetch<'_>>::Item, QuerySingleError> {
+    pub fn single_mut(&mut self) -> Result<<Q::Fetch as Fetch<'_, '_>>::Item, QuerySingleError> {
         let mut query = self.iter_mut();
         let first = query.next();
         let extra = query.next().is_some();