Add query reborrowing (#14690)

# Objective

- Sometimes some method or function takes an owned `Query`, but we don't
want to give up ours;
- transmuting it technically a solution, but it more costly than
necessary.
- Make query iterators more flexible
- this would allow the equivalent of
`slice::split_first`/`slice::split_first_mut` for query iterators
  - helps with requests like #14685

## Solution

- Add a way for reborrowing queries, that is going from a `&'a mut
Query<'w, 's, D, F>` to a `Query<'a, 's, D, F>`:
- this is safe because the original query will be borrowed while the new
query exists and thus no aliased access can happen;
- it's basically the equivalent of going from `&'short mut &'long mut T`
to `&'short mut T` the the compiler automatically implements.
- Add a way for getting the remainder of a query iterator:
- this is interesting also because the original iterator keeps its
position, which was not possible before;
- this in turn requires a way to reborrow query fetches, which I had to
add to `WorldQuery`.

## Showcase

- You can now reborrow a `Query`, getting an equivalent `Query` with a
shorter lifetime. Previously this was possible for read-only queries by
using `Query::to_readonly`, now it's possible for mutable queries too;
- You can now separately iterate over the remainder of `QueryIter`.

## Migration Guide

- `WorldQuery` now has an additional `shrink_fetch` method you have to
implement if you were implementing `WorldQuery` manually.

crates/bevy_ecs/macros/src/world_query.rs

@@ -106,6 +106,17 @@ pub(crate) fn world_query_impl(
                 }
             }
 
+            fn shrink_fetch<'__wlong: '__wshort, '__wshort>(
+                fetch: <#struct_name #user_ty_generics as #path::query::WorldQuery>::Fetch<'__wlong>
+            ) -> <#struct_name #user_ty_generics as #path::query::WorldQuery>::Fetch<'__wshort> {
+                #fetch_struct_name {
+                    #(
+                        #named_field_idents: <#field_types>::shrink_fetch(fetch.#named_field_idents),
+                    )*
+                    #marker_name: &(),
+                }
+            }
+
             unsafe fn init_fetch<'__w>(
                 _world: #path::world::unsafe_world_cell::UnsafeWorldCell<'__w>,
                 state: &Self::State,

crates/bevy_ecs/src/query/fetch.rs

@@ -301,6 +301,8 @@ unsafe impl WorldQuery for Entity {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
+
     unsafe fn init_fetch<'w>(
         _world: UnsafeWorldCell<'w>,
         _state: &Self::State,
@@ -369,6 +371,10 @@ unsafe impl WorldQuery for EntityLocation {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
         _state: &Self::State,
@@ -442,6 +448,10 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
         _state: &Self::State,
@@ -518,6 +528,10 @@ unsafe impl<'a> WorldQuery for EntityMut<'a> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
         _state: &Self::State,
@@ -591,6 +605,10 @@ unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     const IS_DENSE: bool = false;
 
     unsafe fn init_fetch<'w>(
@@ -694,6 +712,10 @@ unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     const IS_DENSE: bool = false;
 
     unsafe fn init_fetch<'w>(
@@ -805,6 +827,10 @@ unsafe impl WorldQuery for &Archetype {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
         _state: &Self::State,
@@ -897,6 +923,10 @@ unsafe impl<T: Component> WorldQuery for &T {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
@@ -1056,6 +1086,10 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
@@ -1251,6 +1285,10 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
@@ -1425,6 +1463,10 @@ unsafe impl<'__w, T: Component> WorldQuery for Mut<'__w, T> {
         <&mut T as WorldQuery>::shrink(item)
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     // Forwarded to `&mut T`
     unsafe fn init_fetch<'w>(
@@ -1535,6 +1577,13 @@ unsafe impl<T: WorldQuery> WorldQuery for Option<T> {
         item.map(T::shrink)
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        OptionFetch {
+            fetch: T::shrink_fetch(fetch.fetch),
+            matches: fetch.matches,
+        }
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
@@ -1712,6 +1761,10 @@ unsafe impl<T: Component> WorldQuery for Has<T> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         _world: UnsafeWorldCell<'w>,
@@ -1830,6 +1883,12 @@ macro_rules! impl_anytuple_fetch {
                     $name.map($name::shrink),
                 )*)
             }
+            fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+                let ($($name,)*) = fetch;
+                ($(
+                    ($name::shrink_fetch($name.0), $name.1),
+                )*)
+            }
 
             #[inline]
             #[allow(clippy::unused_unit)]
@@ -1964,6 +2023,8 @@ unsafe impl<D: QueryData> WorldQuery for NopWorldQuery<D> {
 
     fn shrink<'wlong: 'wshort, 'wshort>(_: ()) {}
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: ()) {}
+
     #[inline(always)]
     unsafe fn init_fetch(
         _world: UnsafeWorldCell,
@@ -2032,6 +2093,9 @@ unsafe impl<T: ?Sized> WorldQuery for PhantomData<T> {
 
     fn shrink<'wlong: 'wshort, 'wshort>(_item: Self::Item<'wlong>) -> Self::Item<'wshort> {}
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(_fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+    }
+
     unsafe fn init_fetch<'w>(
         _world: UnsafeWorldCell<'w>,
         _state: &Self::State,

crates/bevy_ecs/src/query/filter.rs

@@ -142,6 +142,8 @@ unsafe impl<T: Component> WorldQuery for With<T> {
 
     fn shrink<'wlong: 'wshort, 'wshort>(_: Self::Item<'wlong>) -> Self::Item<'wshort> {}
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
+
     #[inline]
     unsafe fn init_fetch(
         _world: UnsafeWorldCell,
@@ -250,6 +252,8 @@ unsafe impl<T: Component> WorldQuery for Without<T> {
 
     fn shrink<'wlong: 'wshort, 'wshort>(_: Self::Item<'wlong>) -> Self::Item<'wshort> {}
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
+
     #[inline]
     unsafe fn init_fetch(
         _world: UnsafeWorldCell,
@@ -386,6 +390,16 @@ macro_rules! impl_or_query_filter {
                 item
             }
 
+            fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+                let ($($filter,)*) = fetch;
+                ($(
+                    OrFetch {
+                        fetch: $filter::shrink_fetch($filter.fetch),
+                        matches: $filter.matches
+                    },
+                )*)
+            }
+
             const IS_DENSE: bool = true $(&& $filter::IS_DENSE)*;
 
             #[inline]
@@ -614,6 +628,10 @@ unsafe impl<T: Component> WorldQuery for Added<T> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,
@@ -825,6 +843,10 @@ unsafe impl<T: Component> WorldQuery for Changed<T> {
         item
     }
 
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+        fetch
+    }
+
     #[inline]
     unsafe fn init_fetch<'w>(
         world: UnsafeWorldCell<'w>,

crates/bevy_ecs/src/query/iter.rs

@@ -50,6 +50,78 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
         }
     }
 
+    /// Creates a new separate iterator yielding the same remaining items of the current one.
+    /// Advancing the new iterator will not advance the original one, which will resume at the
+    /// point it was left at.
+    ///
+    /// Differently from [`remaining_mut`](QueryIter::remaining_mut) the new iterator does not
+    /// borrow from the original one. However it can only be called from an iterator over read only
+    /// items.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # #[derive(Component)]
+    /// # struct ComponentA;
+    ///
+    /// fn combinations(query: Query<&ComponentA>) {
+    ///     let mut iter = query.iter();
+    ///     while let Some(a) = iter.next() {
+    ///         for b in iter.remaining() {
+    ///             // Check every combination (a, b)
+    ///         }
+    ///     }
+    /// }
+    /// ```
+    pub fn remaining(&self) -> QueryIter<'w, 's, D, F>
+    where
+        D: ReadOnlyQueryData,
+    {
+        QueryIter {
+            world: self.world,
+            tables: self.tables,
+            archetypes: self.archetypes,
+            query_state: self.query_state,
+            cursor: self.cursor.clone(),
+        }
+    }
+
+    /// Creates a new separate iterator yielding the same remaining items of the current one.
+    /// Advancing the new iterator will not advance the original one, which will resume at the
+    /// point it was left at.
+    ///
+    /// This method can be called on iterators over mutable items. However the original iterator
+    /// will be borrowed while the new iterator exists and will thus not be usable in that timespan.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # #[derive(Component)]
+    /// # struct ComponentA;
+    ///
+    /// fn combinations(mut query: Query<&mut ComponentA>) {
+    ///     let mut iter = query.iter_mut();
+    ///     while let Some(a) = iter.next() {
+    ///         for b in iter.remaining_mut() {
+    ///             // Check every combination (a, b)
+    ///         }
+    ///     }
+    /// }
+    /// ```
+    pub fn remaining_mut(&mut self) -> QueryIter<'_, 's, D, F> {
+        QueryIter {
+            world: self.world,
+            tables: self.tables,
+            archetypes: self.archetypes,
+            query_state: self.query_state,
+            cursor: self.cursor.reborrow(),
+        }
+    }
+
     /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
     /// from an table.
     ///
@@ -1665,6 +1737,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
         }
     }
 
+    fn reborrow(&mut self) -> QueryIterationCursor<'_, 's, D, F> {
+        QueryIterationCursor {
+            fetch: D::shrink_fetch(self.fetch.clone()),
+            filter: F::shrink_fetch(self.filter.clone()),
+            table_entities: self.table_entities,
+            archetype_entities: self.archetype_entities,
+            storage_id_iter: self.storage_id_iter.clone(),
+            current_len: self.current_len,
+            current_row: self.current_row,
+        }
+    }
+
     /// retrieve item returned from most recent `next` call again.
     #[inline]
     unsafe fn peek_last(&mut self) -> Option<D::Item<'w>> {

crates/bevy_ecs/src/query/world_query.rs

@@ -53,6 +53,9 @@ pub unsafe trait WorldQuery {
     /// This function manually implements subtyping for the query items.
     fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort>;
 
+    /// This function manually implements subtyping for the query fetches.
+    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort>;
+
     /// Creates a new instance of this fetch.
     ///
     /// # Safety
@@ -166,6 +169,13 @@ macro_rules! impl_tuple_world_query {
                 )*)
             }
 
+            fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
+                let ($($name,)*) = fetch;
+                ($(
+                    $name::shrink_fetch($name),
+                )*)
+            }
+
             #[inline]
             #[allow(clippy::unused_unit)]
             unsafe fn init_fetch<'w>(_world: UnsafeWorldCell<'w>, state: &Self::State, _last_run: Tick, _this_run: Tick) -> Self::Fetch<'w> {

crates/bevy_ecs/src/system/query.rs

@@ -405,6 +405,36 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
         unsafe { Query::new(self.world, new_state, self.last_run, self.this_run) }
     }
 
+    /// Returns a new `Query` reborrowing the access from this one. The current query will be unusable
+    /// while the new one exists.
+    ///
+    /// # Example
+    ///
+    /// For example this allows to call other methods or other systems that require an owned `Query` without
+    /// completely giving up ownership of it.
+    ///
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # #[derive(Component)]
+    /// # struct ComponentA;
+    ///
+    /// fn helper_system(query: Query<&ComponentA>) { /* ... */}
+    ///
+    /// fn system(mut query: Query<&ComponentA>) {
+    ///     helper_system(query.reborrow());
+    ///     // Can still use query here:
+    ///     for component in &query {
+    ///         // ...
+    ///     }
+    /// }
+    /// ```
+    pub fn reborrow(&mut self) -> Query<'_, 's, D, F> {
+        // SAFETY: this query is exclusively borrowed while the new one exists, so
+        // no overlapping access can occur.
+        unsafe { Query::new(self.world, self.state, self.last_run, self.this_run) }
+    }
+
     /// Returns an [`Iterator`] over the read-only query items.
     ///
     /// This iterator is always guaranteed to return results from each matching entity once and only once.