Skip to content

Commit 01474f0

Browse files
nicopapnicopap
committed
Improve QueryIter size_hint hints (#4244)
## Objective This fixes #1686. `size_hint` can be useful even if a little niche. For example, `collect::<Vec<_>>()` uses the `size_hint` of Iterator it collects from to pre-allocate a memory slice large enough to not require re-allocating when pushing all the elements of the iterator. ## Solution To this effect I made the following changes: * Add a `IS_ARCHETYPAL` associated constant to the `Fetch` trait, this constant tells us when it is safe to assume that the `Fetch` relies exclusively on archetypes to filter queried entities * Add `IS_ARCHETYPAL` to all the implementations of `Fetch` * Use that constant in `QueryIter::size_hint` to provide a more useful ## Migration guide The new associated constant is an API breaking change. For the user, if they implemented a custom `Fetch`, it means they have to add this associated constant to their implementation. Either `true` if it doesn't limit the number of entities returned in a query beyond that of archetypes, or `false` for when it does.
1 parent 24b2bed commit 01474f0

File tree

6 files changed

+95
-24
lines changed

6 files changed

+95
-24
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,8 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
252252
}
253253
}
254254

255+
const IS_ARCHETYPAL: bool = true #(&& <#field_types as #path::query::WorldQuery>::ReadOnlyFetch::IS_ARCHETYPAL)*;
256+
255257
const IS_DENSE: bool = true #(&& <#field_types as #path::query::WorldQuery>::ReadOnlyFetch::IS_DENSE)*;
256258

257259
#[inline]
@@ -303,6 +305,8 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
303305
}
304306
}
305307

308+
const IS_ARCHETYPAL: bool = true #(&& <#field_types as #path::query::WorldQuery>::#fetch_associated_type::IS_ARCHETYPAL)*;
309+
306310
const IS_DENSE: bool = true #(&& <#field_types as #path::query::WorldQuery>::#fetch_associated_type::IS_DENSE)*;
307311

308312
/// SAFETY: we call `set_archetype` for each member that implements `Fetch`

‎crates/bevy_ecs/src/lib.rs

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,8 @@ mod tests {
5151
component::{Component, ComponentId},
5252
entity::Entity,
5353
query::{
54-
Added, ChangeTrackers, Changed, FilterFetch, FilteredAccess, With, Without, WorldQuery,
54+
Added, ChangeTrackers, Changed, FilterFetch, FilteredAccess, Or, With, Without,
55+
WorldQuery,
5556
},
5657
world::{Mut, World},
5758
};
@@ -1402,6 +1403,40 @@ mod tests {
14021403
);
14031404
}
14041405

1406+
#[test]
1407+
fn test_is_archetypal_size_hints() {
1408+
let mut world = World::default();
1409+
macro_rules! query_min_size {
1410+
($query:ty, $filter:ty) => {
1411+
world
1412+
.query_filtered::<$query, $filter>()
1413+
.iter(&world)
1414+
.size_hint()
1415+
.0
1416+
};
1417+
}
1418+
1419+
world.spawn().insert_bundle((A(1), B(1), C));
1420+
world.spawn().insert_bundle((A(1), C));
1421+
world.spawn().insert_bundle((A(1), B(1)));
1422+
world.spawn().insert_bundle((B(1), C));
1423+
world.spawn().insert(A(1));
1424+
world.spawn().insert(C);
1425+
assert_eq!(2, query_min_size![(), (With<A>, Without<B>)],);
1426+
assert_eq!(3, query_min_size![&B, Or<(With<A>, With<C>)>],);
1427+
assert_eq!(1, query_min_size![&B, (With<A>, With<C>)],);
1428+
assert_eq!(1, query_min_size![(&A, &B), With<C>],);
1429+
assert_eq!(4, query_min_size![&A, ()], "Simple Archetypal");
1430+
assert_eq!(4, query_min_size![ChangeTrackers<A>, ()],);
1431+
// All the following should set minimum size to 0, as it's impossible to predict
1432+
// how many entites the filters will trim.
1433+
assert_eq!(0, query_min_size![(), Added<A>], "Simple Added");
1434+
assert_eq!(0, query_min_size![(), Changed<A>], "Simple Changed");
1435+
assert_eq!(0, query_min_size![(&A, &B), Changed<A>],);
1436+
assert_eq!(0, query_min_size![&A, (Changed<A>, With<B>)],);
1437+
assert_eq!(0, query_min_size![(&A, &B), Or<(Changed<A>, Changed<B>)>],);
1438+
}
1439+
14051440
#[test]
14061441
fn reserve_entities_across_worlds() {
14071442
let mut world_a = World::default();

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,13 @@ pub trait Fetch<'world, 'state>: Sized {
340340
/// [`Fetch::archetype_fetch`] will be called for iterators.
341341
const IS_DENSE: bool;
342342

343+
/// Returns true if (and only if) this Fetch relies strictly on archetypes to limit which
344+
/// components are acessed by the Query.
345+
///
346+
/// This enables optimizations for [`crate::query::QueryIter`] that rely on knowing exactly how
347+
/// many elements are being iterated (such as `Iterator::collect()`).
348+
const IS_ARCHETYPAL: bool;
349+
343350
/// Adjusts internal state to account for the next [`Archetype`]. This will always be called on
344351
/// archetypes that match this [`Fetch`].
345352
///
@@ -458,6 +465,8 @@ impl<'w, 's> Fetch<'w, 's> for EntityFetch {
458465

459466
const IS_DENSE: bool = true;
460467

468+
const IS_ARCHETYPAL: bool = true;
469+
461470
unsafe fn init(
462471
_world: &World,
463472
_state: &Self::State,
@@ -583,6 +592,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
583592
}
584593
};
585594

595+
const IS_ARCHETYPAL: bool = true;
596+
586597
unsafe fn init(
587598
world: &World,
588599
state: &Self::State,
@@ -767,6 +778,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
767778
}
768779
};
769780

781+
const IS_ARCHETYPAL: bool = true;
782+
770783
unsafe fn init(
771784
world: &World,
772785
state: &Self::State,
@@ -873,6 +886,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadOnlyWriteFetch<T> {
873886
}
874887
};
875888

889+
const IS_ARCHETYPAL: bool = true;
890+
876891
unsafe fn init(
877892
world: &World,
878893
state: &Self::State,
@@ -1002,6 +1017,8 @@ impl<'w, 's, T: Fetch<'w, 's>> Fetch<'w, 's> for OptionFetch<T> {
10021017

10031018
const IS_DENSE: bool = T::IS_DENSE;
10041019

1020+
const IS_ARCHETYPAL: bool = T::IS_ARCHETYPAL;
1021+
10051022
unsafe fn init(
10061023
world: &World,
10071024
state: &Self::State,
@@ -1212,6 +1229,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ChangeTrackersFetch<T> {
12121229
}
12131230
};
12141231

1232+
const IS_ARCHETYPAL: bool = true;
1233+
12151234
unsafe fn init(
12161235
world: &World,
12171236
state: &Self::State,
@@ -1314,6 +1333,8 @@ macro_rules! impl_tuple_fetch {
13141333

13151334
const IS_DENSE: bool = true $(&& $name::IS_DENSE)*;
13161335

1336+
const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*;
1337+
13171338
#[inline]
13181339
unsafe fn set_archetype(&mut self, _state: &Self::State, _archetype: &Archetype, _tables: &Tables) {
13191340
let ($($name,)*) = self;
@@ -1407,6 +1428,8 @@ macro_rules! impl_anytuple_fetch {
14071428

14081429
const IS_DENSE: bool = true $(&& $name::IS_DENSE)*;
14091430

1431+
const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*;
1432+
14101433
#[inline]
14111434
unsafe fn set_archetype(&mut self, _state: &Self::State, _archetype: &Archetype, _tables: &Tables) {
14121435
let ($($name,)*) = &mut self.0;
@@ -1512,6 +1535,8 @@ impl<'w, 's, State: FetchState> Fetch<'w, 's> for NopFetch<State> {
15121535

15131536
const IS_DENSE: bool = true;
15141537

1538+
const IS_ARCHETYPAL: bool = true;
1539+
15151540
#[inline(always)]
15161541
unsafe fn init(
15171542
_world: &World,

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch<T> {
145145
}
146146
};
147147

148+
const IS_ARCHETYPAL: bool = true;
149+
148150
#[inline]
149151
unsafe fn set_table(&mut self, _state: &Self::State, _table: &Table) {}
150152

@@ -280,6 +282,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithoutFetch<T> {
280282
}
281283
};
282284

285+
const IS_ARCHETYPAL: bool = true;
286+
283287
#[inline]
284288
unsafe fn set_table(&mut self, _state: &Self::State, _table: &Table) {}
285289

@@ -402,6 +406,8 @@ macro_rules! impl_query_filter_tuple {
402406

403407
const IS_DENSE: bool = true $(&& $filter::IS_DENSE)*;
404408

409+
const IS_ARCHETYPAL: bool = true $(&& $filter::IS_ARCHETYPAL)*;
410+
405411
#[inline]
406412
unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
407413
let ($($filter,)*) = &mut self.0;
@@ -578,6 +584,8 @@ macro_rules! impl_tick_filter {
578584
}
579585
};
580586

587+
const IS_ARCHETYPAL: bool = false;
588+
581589
unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
582590
self.table_ticks = table
583591
.get_column(state.component_id).unwrap()

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

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -138,9 +138,6 @@ where
138138
}
139139
}
140140

141-
// NOTE: For unfiltered Queries this should actually return a exact size hint,
142-
// to fulfil the ExactSizeIterator invariant, but this isn't practical without specialization.
143-
// For more information see Issue #1686.
144141
fn size_hint(&self) -> (usize, Option<usize>) {
145142
let max_size = self
146143
.query_state
@@ -149,7 +146,9 @@ where
149146
.map(|index| self.world.archetypes[ArchetypeId::new(index)].len())
150147
.sum();
151148

152-
(0, Some(max_size))
149+
let archetype_query = F::Fetch::IS_ARCHETYPAL && QF::IS_ARCHETYPAL;
150+
let min_size = if archetype_query { max_size } else { 0 };
151+
(min_size, Some(max_size))
153152
}
154153
}
155154

@@ -297,9 +296,6 @@ where
297296
unsafe { QueryCombinationIter::fetch_next_aliased_unchecked(self) }
298297
}
299298

300-
// NOTE: For unfiltered Queries this should actually return a exact size hint,
301-
// to fulfil the ExactSizeIterator invariant, but this isn't practical without specialization.
302-
// For more information see Issue #1686.
303299
fn size_hint(&self) -> (usize, Option<usize>) {
304300
if K == 0 {
305301
return (0, Some(0));
@@ -324,15 +320,18 @@ where
324320
n / k_factorial
325321
});
326322

327-
(0, max_combinations)
323+
let archetype_query = F::Fetch::IS_ARCHETYPAL && QF::IS_ARCHETYPAL;
324+
let min_combinations = if archetype_query { max_size } else { 0 };
325+
(min_combinations, max_combinations)
328326
}
329327
}
330328

331329
// NOTE: We can cheaply implement this for unfiltered Queries because we have:
332330
// (1) pre-computed archetype matches
333331
// (2) each archetype pre-computes length
334332
// (3) there are no per-entity filters
335-
// TODO: add an ArchetypeOnlyFilter that enables us to implement this for filters like With<T>
333+
// TODO: add an ArchetypeOnlyFilter that enables us to implement this for filters like With<T>.
334+
// This would need to be added to all types that implement Filter with Filter::IS_ARCHETYPAL = true
336335
impl<'w, 's, Q: WorldQuery, QF> ExactSizeIterator for QueryIter<'w, 's, Q, QF, ()>
337336
where
338337
QF: Fetch<'w, 's, State = Q::State>,

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

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -54,38 +54,38 @@ mod tests {
5454
let mut a_query = world.query::<&A>();
5555
assert_eq!(a_query.iter_combinations::<0>(&world).count(), 0);
5656
assert_eq!(
57-
a_query.iter_combinations::<0>(&world).size_hint(),
58-
(0, Some(0))
57+
a_query.iter_combinations::<0>(&world).size_hint().1,
58+
Some(0)
5959
);
6060
assert_eq!(a_query.iter_combinations::<1>(&world).count(), 4);
6161
assert_eq!(
62-
a_query.iter_combinations::<1>(&world).size_hint(),
63-
(0, Some(4))
62+
a_query.iter_combinations::<1>(&world).size_hint().1,
63+
Some(4)
6464
);
6565
assert_eq!(a_query.iter_combinations::<2>(&world).count(), 6);
6666
assert_eq!(
67-
a_query.iter_combinations::<2>(&world).size_hint(),
68-
(0, Some(6))
67+
a_query.iter_combinations::<2>(&world).size_hint().1,
68+
Some(6)
6969
);
7070
assert_eq!(a_query.iter_combinations::<3>(&world).count(), 4);
7171
assert_eq!(
72-
a_query.iter_combinations::<3>(&world).size_hint(),
73-
(0, Some(4))
72+
a_query.iter_combinations::<3>(&world).size_hint().1,
73+
Some(4)
7474
);
7575
assert_eq!(a_query.iter_combinations::<4>(&world).count(), 1);
7676
assert_eq!(
77-
a_query.iter_combinations::<4>(&world).size_hint(),
78-
(0, Some(1))
77+
a_query.iter_combinations::<4>(&world).size_hint().1,
78+
Some(1)
7979
);
8080
assert_eq!(a_query.iter_combinations::<5>(&world).count(), 0);
8181
assert_eq!(
82-
a_query.iter_combinations::<5>(&world).size_hint(),
83-
(0, Some(0))
82+
a_query.iter_combinations::<5>(&world).size_hint().1,
83+
Some(0)
8484
);
8585
assert_eq!(a_query.iter_combinations::<1024>(&world).count(), 0);
8686
assert_eq!(
87-
a_query.iter_combinations::<1024>(&world).size_hint(),
88-
(0, Some(0))
87+
a_query.iter_combinations::<1024>(&world).size_hint().1,
88+
Some(0)
8989
);
9090

9191
let values: Vec<[&A; 2]> = world.query::<&A>().iter_combinations(&world).collect();

0 commit comments

Comments
 (0)