From e9e9e5e15da4528c3d33c0808e954cf7a8b874f7 Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Thu, 15 Aug 2024 19:38:56 +0200 Subject: [PATCH] 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 | 11 +++ crates/bevy_ecs/src/query/fetch.rs | 64 +++++++++++++++++ crates/bevy_ecs/src/query/filter.rs | 22 ++++++ crates/bevy_ecs/src/query/iter.rs | 84 +++++++++++++++++++++++ crates/bevy_ecs/src/query/world_query.rs | 10 +++ crates/bevy_ecs/src/system/query.rs | 30 ++++++++ 6 files changed, 221 insertions(+) diff --git a/crates/bevy_ecs/macros/src/world_query.rs b/crates/bevy_ecs/macros/src/world_query.rs index 42fd8ccb59..26cb8edadd 100644 --- a/crates/bevy_ecs/macros/src/world_query.rs +++ b/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, diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 269b9dc914..8e522b61a2 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/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 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 WorldQuery for Option { 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 WorldQuery for Has { 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 WorldQuery for NopWorldQuery { 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 WorldQuery for PhantomData { 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, diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index 93036a5c5b..affdaf2828 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -142,6 +142,8 @@ unsafe impl WorldQuery for With { 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 WorldQuery for Without { 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 WorldQuery for Added { 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 WorldQuery for Changed { item } + fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { + fetch + } + #[inline] unsafe fn init_fetch<'w>( world: UnsafeWorldCell<'w>, diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index f2c74429d6..f8564025aa 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/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> { diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index 6e7d0479cf..f2da8b3558 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/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> { diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index e270c028d0..ffe595b2cf 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/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.