Skip to content

Commit

Permalink
Improve QueryIter size_hint hints
Browse files Browse the repository at this point in the history
This fixes bevyengine#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.

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

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.
  • Loading branch information
nicopap committed Mar 17, 2022
1 parent c1a2378 commit 36d2534
Show file tree
Hide file tree
Showing 5 changed files with 78 additions and 2 deletions.
4 changes: 4 additions & 0 deletions crates/bevy_ecs/macros/src/fetch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,8 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
}
}

const IS_ARCHETYPAL: bool = true #(&& <#field_types as #path::query::WorldQuery>::ReadOnlyFetch::IS_ARCHETYPAL)*;

const IS_DENSE: bool = true #(&& <#field_types as #path::query::WorldQuery>::ReadOnlyFetch::IS_DENSE)*;

#[inline]
Expand Down Expand Up @@ -298,6 +300,8 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
}
}

const IS_ARCHETYPAL: bool = true #(&& <#field_types as #path::query::WorldQuery>::#fetch_associated_type::IS_ARCHETYPAL)*;

const IS_DENSE: bool = true #(&& <#field_types as #path::query::WorldQuery>::#fetch_associated_type::IS_DENSE)*;

/// SAFETY: we call `set_archetype` for each member that implements `Fetch`
Expand Down
37 changes: 36 additions & 1 deletion crates/bevy_ecs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ mod tests {
component::{Component, ComponentId},
entity::Entity,
query::{
Added, ChangeTrackers, Changed, FilterFetch, FilteredAccess, With, Without, WorldQuery,
Added, ChangeTrackers, Changed, FilterFetch, FilteredAccess, Or, With, Without,
WorldQuery,
},
world::{Mut, World},
};
Expand Down Expand Up @@ -1392,6 +1393,40 @@ mod tests {
);
}

#[test]
fn test_is_archetypal_size_hints() {
let mut world = World::default();
macro_rules! query_min_size {
($query:ty, $filter:ty) => {
world
.query_filtered::<$query, $filter>()
.iter(&world)
.size_hint()
.0
};
}

world.spawn().insert_bundle((A(1), B(1), C));
world.spawn().insert_bundle((A(1), C));
world.spawn().insert_bundle((A(1), B(1)));
world.spawn().insert_bundle((B(1), C));
world.spawn().insert(A(1));
world.spawn().insert(C);
assert_eq!(2, query_min_size![(), (With<A>, Without<B>)],);
assert_eq!(3, query_min_size![&B, Or<(With<A>, With<C>)>],);
assert_eq!(1, query_min_size![&B, (With<A>, With<C>)],);
assert_eq!(1, query_min_size![(&A, &B), With<C>],);
assert_eq!(4, query_min_size![&A, ()], "Simple Archetypal");
assert_eq!(4, query_min_size![ChangeTrackers<A>, ()],);
// All the following should set minimum size to 0, as it's impossible to predict
// how many entites the filters will trim.
assert_eq!(0, query_min_size![(), Added<A>], "Simple Added");
assert_eq!(0, query_min_size![(), Changed<A>], "Simple Changed");
assert_eq!(0, query_min_size![(&A, &B), Changed<A>],);
assert_eq!(0, query_min_size![&A, (Changed<A>, With<B>)],);
assert_eq!(0, query_min_size![(&A, &B), Or<(Changed<A>, Changed<B>)>],);
}

#[test]
fn reserve_entities_across_worlds() {
let mut world_a = World::default();
Expand Down
25 changes: 25 additions & 0 deletions crates/bevy_ecs/src/query/fetch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,13 @@ pub trait Fetch<'world, 'state>: Sized {
/// [`Fetch::archetype_fetch`] will be called for iterators.
const IS_DENSE: bool;

/// Returns true if (and only if) this Fetch relies strictly on archetypes to limit which
/// components are acessed by the Query.
///
/// This enables optimizations for [`Query::iter`] that rely on knowing exactly how many
/// elements are being iterated (such as `Iterator::collect()`).
const IS_ARCHETYPAL: bool;

/// Adjusts internal state to account for the next [`Archetype`]. This will always be called on
/// archetypes that match this [`Fetch`].
///
Expand Down Expand Up @@ -456,6 +463,8 @@ impl<'w, 's> Fetch<'w, 's> for EntityFetch {

const IS_DENSE: bool = true;

const IS_ARCHETYPAL: bool = true;

unsafe fn init(
_world: &World,
_state: &Self::State,
Expand Down Expand Up @@ -579,6 +588,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
}
};

const IS_ARCHETYPAL: bool = true;

unsafe fn init(
world: &World,
state: &Self::State,
Expand Down Expand Up @@ -760,6 +771,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
}
};

const IS_ARCHETYPAL: bool = true;

unsafe fn init(
world: &World,
state: &Self::State,
Expand Down Expand Up @@ -866,6 +879,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadOnlyWriteFetch<T> {
}
};

const IS_ARCHETYPAL: bool = true;

unsafe fn init(
world: &World,
state: &Self::State,
Expand Down Expand Up @@ -993,6 +1008,8 @@ impl<'w, 's, T: Fetch<'w, 's>> Fetch<'w, 's> for OptionFetch<T> {

const IS_DENSE: bool = T::IS_DENSE;

const IS_ARCHETYPAL: bool = T::IS_ARCHETYPAL;

unsafe fn init(
world: &World,
state: &Self::State,
Expand Down Expand Up @@ -1201,6 +1218,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ChangeTrackersFetch<T> {
}
};

const IS_ARCHETYPAL: bool = true;

unsafe fn init(
world: &World,
state: &Self::State,
Expand Down Expand Up @@ -1303,6 +1322,8 @@ macro_rules! impl_tuple_fetch {

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

const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*;

#[inline]
unsafe fn set_archetype(&mut self, _state: &Self::State, _archetype: &Archetype, _tables: &Tables) {
let ($($name,)*) = self;
Expand Down Expand Up @@ -1396,6 +1417,8 @@ macro_rules! impl_anytuple_fetch {

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

const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*;

#[inline]
unsafe fn set_archetype(&mut self, _state: &Self::State, _archetype: &Archetype, _tables: &Tables) {
let ($($name,)*) = &mut self.0;
Expand Down Expand Up @@ -1501,6 +1524,8 @@ impl<'w, 's, State: FetchState> Fetch<'w, 's> for NopFetch<State> {

const IS_DENSE: bool = true;

const IS_ARCHETYPAL: bool = true;

#[inline(always)]
unsafe fn init(
_world: &World,
Expand Down
8 changes: 8 additions & 0 deletions crates/bevy_ecs/src/query/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch<T> {
}
};

const IS_ARCHETYPAL: bool = true;

#[inline]
unsafe fn set_table(&mut self, _state: &Self::State, _table: &Table) {}

Expand Down Expand Up @@ -266,6 +268,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithoutFetch<T> {
}
};

const IS_ARCHETYPAL: bool = true;

#[inline]
unsafe fn set_table(&mut self, _state: &Self::State, _table: &Table) {}

Expand Down Expand Up @@ -375,6 +379,8 @@ macro_rules! impl_query_filter_tuple {

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

const IS_ARCHETYPAL: bool = true $(&& $filter::IS_ARCHETYPAL)*;

#[inline]
unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
let ($($filter,)*) = &mut self.0;
Expand Down Expand Up @@ -549,6 +555,8 @@ macro_rules! impl_tick_filter {
}
};

const IS_ARCHETYPAL: bool = false;

unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
self.table_ticks = table
.get_column(state.component_id).unwrap()
Expand Down
6 changes: 5 additions & 1 deletion crates/bevy_ecs/src/query/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,11 @@ where
.map(|index| self.world.archetypes[ArchetypeId::new(index)].len())
.sum();

(0, Some(max_size))
// TODO: it's _probably possible to use const generics to have specialized implementation
// of size_hint based on whether this is true or not.
let archetype_query = F::Fetch::IS_ARCHETYPAL && QF::IS_ARCHETYPAL;
let min_size = if archetype_query { max_size } else { 0 };
(min_size, Some(max_size))
}
}

Expand Down

0 comments on commit 36d2534

Please sign in to comment.