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> {