Mirrors/bevy
2
0
Fork 0
mirror of https://github.com/bevyengine/bevy synced 2024-11-10 07:04:33 +00:00
Code Issues Projects Releases Packages Wiki Activity

Add query reborrowing (#14690)

Browse source 
# 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.
This commit is contained in:
Giacomo Stevanato 2024-08-15 19:38:56 +02:00 committed by GitHub
parent 9d5837769c
commit e9e9e5e15d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 221 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
crates/bevy_ecs/macros/src/world_query.rs
View file

@ -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,

64
crates/bevy_ecs/src/query/fetch.rs
View file

@ -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,

22
crates/bevy_ecs/src/query/filter.rs
View file

@ -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>,

84
crates/bevy_ecs/src/query/iter.rs
View file

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

10
crates/bevy_ecs/src/query/world_query.rs
View file

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

30
crates/bevy_ecs/src/system/query.rs
View file

@ -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.