Skip to content

Commit e9e9e5e

Browse files
SkiFire13SkiFire13
authored
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.
1 parent 9d58377 commit e9e9e5e

File tree

6 files changed

+221
-0
lines changed

6 files changed

+221
-0
lines changed

‎crates/bevy_ecs/macros/src/world_query.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,17 @@ pub(crate) fn world_query_impl(
106106
}
107107
}
108108

109+
fn shrink_fetch<'__wlong: '__wshort, '__wshort>(
110+
fetch: <#struct_name #user_ty_generics as #path::query::WorldQuery>::Fetch<'__wlong>
111+
) -> <#struct_name #user_ty_generics as #path::query::WorldQuery>::Fetch<'__wshort> {
112+
#fetch_struct_name {
113+
#(
114+
#named_field_idents: <#field_types>::shrink_fetch(fetch.#named_field_idents),
115+
)*
116+
#marker_name: &(),
117+
}
118+
}
119+
109120
unsafe fn init_fetch<'__w>(
110121
_world: #path::world::unsafe_world_cell::UnsafeWorldCell<'__w>,
111122
state: &Self::State,

‎crates/bevy_ecs/src/query/fetch.rs

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,8 @@ unsafe impl WorldQuery for Entity {
301301
item
302302
}
303303

304+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
305+
304306
unsafe fn init_fetch<'w>(
305307
_world: UnsafeWorldCell<'w>,
306308
_state: &Self::State,
@@ -369,6 +371,10 @@ unsafe impl WorldQuery for EntityLocation {
369371
item
370372
}
371373

374+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
375+
fetch
376+
}
377+
372378
unsafe fn init_fetch<'w>(
373379
world: UnsafeWorldCell<'w>,
374380
_state: &Self::State,
@@ -442,6 +448,10 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> {
442448
item
443449
}
444450

451+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
452+
fetch
453+
}
454+
445455
unsafe fn init_fetch<'w>(
446456
world: UnsafeWorldCell<'w>,
447457
_state: &Self::State,
@@ -518,6 +528,10 @@ unsafe impl<'a> WorldQuery for EntityMut<'a> {
518528
item
519529
}
520530

531+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
532+
fetch
533+
}
534+
521535
unsafe fn init_fetch<'w>(
522536
world: UnsafeWorldCell<'w>,
523537
_state: &Self::State,
@@ -591,6 +605,10 @@ unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> {
591605
item
592606
}
593607

608+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
609+
fetch
610+
}
611+
594612
const IS_DENSE: bool = false;
595613

596614
unsafe fn init_fetch<'w>(
@@ -694,6 +712,10 @@ unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> {
694712
item
695713
}
696714

715+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
716+
fetch
717+
}
718+
697719
const IS_DENSE: bool = false;
698720

699721
unsafe fn init_fetch<'w>(
@@ -805,6 +827,10 @@ unsafe impl WorldQuery for &Archetype {
805827
item
806828
}
807829

830+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
831+
fetch
832+
}
833+
808834
unsafe fn init_fetch<'w>(
809835
world: UnsafeWorldCell<'w>,
810836
_state: &Self::State,
@@ -897,6 +923,10 @@ unsafe impl<T: Component> WorldQuery for &T {
897923
item
898924
}
899925

926+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
927+
fetch
928+
}
929+
900930
#[inline]
901931
unsafe fn init_fetch<'w>(
902932
world: UnsafeWorldCell<'w>,
@@ -1056,6 +1086,10 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
10561086
item
10571087
}
10581088

1089+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
1090+
fetch
1091+
}
1092+
10591093
#[inline]
10601094
unsafe fn init_fetch<'w>(
10611095
world: UnsafeWorldCell<'w>,
@@ -1251,6 +1285,10 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
12511285
item
12521286
}
12531287

1288+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
1289+
fetch
1290+
}
1291+
12541292
#[inline]
12551293
unsafe fn init_fetch<'w>(
12561294
world: UnsafeWorldCell<'w>,
@@ -1425,6 +1463,10 @@ unsafe impl<'__w, T: Component> WorldQuery for Mut<'__w, T> {
14251463
<&mut T as WorldQuery>::shrink(item)
14261464
}
14271465

1466+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
1467+
fetch
1468+
}
1469+
14281470
#[inline]
14291471
// Forwarded to `&mut T`
14301472
unsafe fn init_fetch<'w>(
@@ -1535,6 +1577,13 @@ unsafe impl<T: WorldQuery> WorldQuery for Option<T> {
15351577
item.map(T::shrink)
15361578
}
15371579

1580+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
1581+
OptionFetch {
1582+
fetch: T::shrink_fetch(fetch.fetch),
1583+
matches: fetch.matches,
1584+
}
1585+
}
1586+
15381587
#[inline]
15391588
unsafe fn init_fetch<'w>(
15401589
world: UnsafeWorldCell<'w>,
@@ -1712,6 +1761,10 @@ unsafe impl<T: Component> WorldQuery for Has<T> {
17121761
item
17131762
}
17141763

1764+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
1765+
fetch
1766+
}
1767+
17151768
#[inline]
17161769
unsafe fn init_fetch<'w>(
17171770
_world: UnsafeWorldCell<'w>,
@@ -1830,6 +1883,12 @@ macro_rules! impl_anytuple_fetch {
18301883
$name.map($name::shrink),
18311884
)*)
18321885
}
1886+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
1887+
let ($($name,)*) = fetch;
1888+
($(
1889+
($name::shrink_fetch($name.0), $name.1),
1890+
)*)
1891+
}
18331892

18341893
#[inline]
18351894
#[allow(clippy::unused_unit)]
@@ -1964,6 +2023,8 @@ unsafe impl<D: QueryData> WorldQuery for NopWorldQuery<D> {
19642023

19652024
fn shrink<'wlong: 'wshort, 'wshort>(_: ()) {}
19662025

2026+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: ()) {}
2027+
19672028
#[inline(always)]
19682029
unsafe fn init_fetch(
19692030
_world: UnsafeWorldCell,
@@ -2032,6 +2093,9 @@ unsafe impl<T: ?Sized> WorldQuery for PhantomData<T> {
20322093

20332094
fn shrink<'wlong: 'wshort, 'wshort>(_item: Self::Item<'wlong>) -> Self::Item<'wshort> {}
20342095

2096+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(_fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
2097+
}
2098+
20352099
unsafe fn init_fetch<'w>(
20362100
_world: UnsafeWorldCell<'w>,
20372101
_state: &Self::State,

‎crates/bevy_ecs/src/query/filter.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,8 @@ unsafe impl<T: Component> WorldQuery for With<T> {
142142

143143
fn shrink<'wlong: 'wshort, 'wshort>(_: Self::Item<'wlong>) -> Self::Item<'wshort> {}
144144

145+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
146+
145147
#[inline]
146148
unsafe fn init_fetch(
147149
_world: UnsafeWorldCell,
@@ -250,6 +252,8 @@ unsafe impl<T: Component> WorldQuery for Without<T> {
250252

251253
fn shrink<'wlong: 'wshort, 'wshort>(_: Self::Item<'wlong>) -> Self::Item<'wshort> {}
252254

255+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(_: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {}
256+
253257
#[inline]
254258
unsafe fn init_fetch(
255259
_world: UnsafeWorldCell,
@@ -386,6 +390,16 @@ macro_rules! impl_or_query_filter {
386390
item
387391
}
388392

393+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
394+
let ($($filter,)*) = fetch;
395+
($(
396+
OrFetch {
397+
fetch: $filter::shrink_fetch($filter.fetch),
398+
matches: $filter.matches
399+
},
400+
)*)
401+
}
402+
389403
const IS_DENSE: bool = true $(&& $filter::IS_DENSE)*;
390404

391405
#[inline]
@@ -614,6 +628,10 @@ unsafe impl<T: Component> WorldQuery for Added<T> {
614628
item
615629
}
616630

631+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
632+
fetch
633+
}
634+
617635
#[inline]
618636
unsafe fn init_fetch<'w>(
619637
world: UnsafeWorldCell<'w>,
@@ -825,6 +843,10 @@ unsafe impl<T: Component> WorldQuery for Changed<T> {
825843
item
826844
}
827845

846+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
847+
fetch
848+
}
849+
828850
#[inline]
829851
unsafe fn init_fetch<'w>(
830852
world: UnsafeWorldCell<'w>,

‎crates/bevy_ecs/src/query/iter.rs

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,78 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
5050
}
5151
}
5252

53+
/// Creates a new separate iterator yielding the same remaining items of the current one.
54+
/// Advancing the new iterator will not advance the original one, which will resume at the
55+
/// point it was left at.
56+
///
57+
/// Differently from [`remaining_mut`](QueryIter::remaining_mut) the new iterator does not
58+
/// borrow from the original one. However it can only be called from an iterator over read only
59+
/// items.
60+
///
61+
/// # Example
62+
///
63+
/// ```
64+
/// # use bevy_ecs::prelude::*;
65+
/// #
66+
/// # #[derive(Component)]
67+
/// # struct ComponentA;
68+
///
69+
/// fn combinations(query: Query<&ComponentA>) {
70+
/// let mut iter = query.iter();
71+
/// while let Some(a) = iter.next() {
72+
/// for b in iter.remaining() {
73+
/// // Check every combination (a, b)
74+
/// }
75+
/// }
76+
/// }
77+
/// ```
78+
pub fn remaining(&self) -> QueryIter<'w, 's, D, F>
79+
where
80+
D: ReadOnlyQueryData,
81+
{
82+
QueryIter {
83+
world: self.world,
84+
tables: self.tables,
85+
archetypes: self.archetypes,
86+
query_state: self.query_state,
87+
cursor: self.cursor.clone(),
88+
}
89+
}
90+
91+
/// Creates a new separate iterator yielding the same remaining items of the current one.
92+
/// Advancing the new iterator will not advance the original one, which will resume at the
93+
/// point it was left at.
94+
///
95+
/// This method can be called on iterators over mutable items. However the original iterator
96+
/// will be borrowed while the new iterator exists and will thus not be usable in that timespan.
97+
///
98+
/// # Example
99+
///
100+
/// ```
101+
/// # use bevy_ecs::prelude::*;
102+
/// #
103+
/// # #[derive(Component)]
104+
/// # struct ComponentA;
105+
///
106+
/// fn combinations(mut query: Query<&mut ComponentA>) {
107+
/// let mut iter = query.iter_mut();
108+
/// while let Some(a) = iter.next() {
109+
/// for b in iter.remaining_mut() {
110+
/// // Check every combination (a, b)
111+
/// }
112+
/// }
113+
/// }
114+
/// ```
115+
pub fn remaining_mut(&mut self) -> QueryIter<'_, 's, D, F> {
116+
QueryIter {
117+
world: self.world,
118+
tables: self.tables,
119+
archetypes: self.archetypes,
120+
query_state: self.query_state,
121+
cursor: self.cursor.reborrow(),
122+
}
123+
}
124+
53125
/// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
54126
/// from an table.
55127
///
@@ -1665,6 +1737,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
16651737
}
16661738
}
16671739

1740+
fn reborrow(&mut self) -> QueryIterationCursor<'_, 's, D, F> {
1741+
QueryIterationCursor {
1742+
fetch: D::shrink_fetch(self.fetch.clone()),
1743+
filter: F::shrink_fetch(self.filter.clone()),
1744+
table_entities: self.table_entities,
1745+
archetype_entities: self.archetype_entities,
1746+
storage_id_iter: self.storage_id_iter.clone(),
1747+
current_len: self.current_len,
1748+
current_row: self.current_row,
1749+
}
1750+
}
1751+
16681752
/// retrieve item returned from most recent `next` call again.
16691753
#[inline]
16701754
unsafe fn peek_last(&mut self) -> Option<D::Item<'w>> {

‎crates/bevy_ecs/src/query/world_query.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ pub unsafe trait WorldQuery {
5353
/// This function manually implements subtyping for the query items.
5454
fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort>;
5555

56+
/// This function manually implements subtyping for the query fetches.
57+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort>;
58+
5659
/// Creates a new instance of this fetch.
5760
///
5861
/// # Safety
@@ -166,6 +169,13 @@ macro_rules! impl_tuple_world_query {
166169
)*)
167170
}
168171

172+
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
173+
let ($($name,)*) = fetch;
174+
($(
175+
$name::shrink_fetch($name),
176+
)*)
177+
}
178+
169179
#[inline]
170180
#[allow(clippy::unused_unit)]
171181
unsafe fn init_fetch<'w>(_world: UnsafeWorldCell<'w>, state: &Self::State, _last_run: Tick, _this_run: Tick) -> Self::Fetch<'w> {

0 commit comments

Comments
 (0)