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.
This commit is contained in:
Nicola Papale 2022-04-27 18:02:06 +00:00
parent 3d4e0066f4
commit 71a246ce9e
6 changed files with 131 additions and 125 deletions

View file

@ -252,6 +252,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]
@ -303,6 +305,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`

View file

@ -51,7 +51,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},
};
@ -1402,6 +1403,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();

View file

@ -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 [`crate::query::QueryIter`] 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`].
///
@ -458,6 +465,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,
@ -583,6 +592,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,
@ -767,6 +778,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,
@ -873,6 +886,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,
@ -1002,6 +1017,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,
@ -1212,6 +1229,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,
@ -1314,6 +1333,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;
@ -1407,6 +1428,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;
@ -1512,6 +1535,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,

View file

@ -145,6 +145,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) {}
@ -280,6 +282,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) {}
@ -402,6 +406,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;
@ -578,6 +584,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()

View file

@ -138,9 +138,6 @@ where
}
}
// NOTE: For unfiltered Queries this should actually return a exact size hint,
// to fulfil the ExactSizeIterator invariant, but this isn't practical without specialization.
// For more information see Issue #1686.
fn size_hint(&self) -> (usize, Option<usize>) {
let max_size = self
.query_state
@ -149,7 +146,9 @@ where
.map(|index| self.world.archetypes[ArchetypeId::new(index)].len())
.sum();
(0, Some(max_size))
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))
}
}
@ -297,9 +296,6 @@ where
unsafe { QueryCombinationIter::fetch_next_aliased_unchecked(self) }
}
// NOTE: For unfiltered Queries this should actually return a exact size hint,
// to fulfil the ExactSizeIterator invariant, but this isn't practical without specialization.
// For more information see Issue #1686.
fn size_hint(&self) -> (usize, Option<usize>) {
if K == 0 {
return (0, Some(0));
@ -324,7 +320,9 @@ where
n / k_factorial
});
(0, max_combinations)
let archetype_query = F::Fetch::IS_ARCHETYPAL && QF::IS_ARCHETYPAL;
let min_combinations = if archetype_query { max_size } else { 0 };
(min_combinations, max_combinations)
}
}
@ -332,7 +330,8 @@ where
// (1) pre-computed archetype matches
// (2) each archetype pre-computes length
// (3) there are no per-entity filters
// TODO: add an ArchetypeOnlyFilter that enables us to implement this for filters like With<T>
// TODO: add an ArchetypeOnlyFilter that enables us to implement this for filters like With<T>.
// This would need to be added to all types that implement Filter with Filter::IS_ARCHETYPAL = true
impl<'w, 's, Q: WorldQuery, QF> ExactSizeIterator for QueryIter<'w, 's, Q, QF, ()>
where
QF: Fetch<'w, 's, State = Q::State>,

View file

@ -52,41 +52,21 @@ mod tests {
world.spawn().insert_bundle((A(4),));
let mut a_query = world.query::<&A>();
assert_eq!(a_query.iter_combinations::<0>(&world).count(), 0);
assert_eq!(
a_query.iter_combinations::<0>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query.iter_combinations::<1>(&world).count(), 4);
assert_eq!(
a_query.iter_combinations::<1>(&world).size_hint(),
(0, Some(4))
);
assert_eq!(a_query.iter_combinations::<2>(&world).count(), 6);
assert_eq!(
a_query.iter_combinations::<2>(&world).size_hint(),
(0, Some(6))
);
assert_eq!(a_query.iter_combinations::<3>(&world).count(), 4);
assert_eq!(
a_query.iter_combinations::<3>(&world).size_hint(),
(0, Some(4))
);
assert_eq!(a_query.iter_combinations::<4>(&world).count(), 1);
assert_eq!(
a_query.iter_combinations::<4>(&world).size_hint(),
(0, Some(1))
);
assert_eq!(a_query.iter_combinations::<5>(&world).count(), 0);
assert_eq!(
a_query.iter_combinations::<5>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query.iter_combinations::<1024>(&world).count(), 0);
assert_eq!(
a_query.iter_combinations::<1024>(&world).size_hint(),
(0, Some(0))
);
let w = &world;
assert_eq!(a_query.iter_combinations::<0>(w).count(), 0);
assert_eq!(a_query.iter_combinations::<0>(w).size_hint().1, Some(0));
assert_eq!(a_query.iter_combinations::<1>(w).count(), 4);
assert_eq!(a_query.iter_combinations::<1>(w).size_hint().1, Some(4));
assert_eq!(a_query.iter_combinations::<2>(w).count(), 6);
assert_eq!(a_query.iter_combinations::<2>(w).size_hint().1, Some(6));
assert_eq!(a_query.iter_combinations::<3>(w).count(), 4);
assert_eq!(a_query.iter_combinations::<3>(w).size_hint().1, Some(4));
assert_eq!(a_query.iter_combinations::<4>(w).count(), 1);
assert_eq!(a_query.iter_combinations::<4>(w).size_hint().1, Some(1));
assert_eq!(a_query.iter_combinations::<5>(w).count(), 0);
assert_eq!(a_query.iter_combinations::<5>(w).size_hint().1, Some(0));
assert_eq!(a_query.iter_combinations::<1024>(w).count(), 0);
assert_eq!(a_query.iter_combinations::<1024>(w).size_hint().1, Some(0));
let values: Vec<[&A; 2]> = world.query::<&A>().iter_combinations(&world).collect();
assert_eq!(
@ -152,86 +132,41 @@ mod tests {
world.spawn().insert_bundle((A(3),));
world.spawn().insert_bundle((A(4),));
let mut a_query_with_b = world.query_filtered::<&A, With<B>>();
assert_eq!(a_query_with_b.iter_combinations::<0>(&world).count(), 0);
assert_eq!(
a_query_with_b.iter_combinations::<0>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_with_b.iter_combinations::<1>(&world).count(), 1);
assert_eq!(
a_query_with_b.iter_combinations::<1>(&world).size_hint(),
(0, Some(1))
);
assert_eq!(a_query_with_b.iter_combinations::<2>(&world).count(), 0);
assert_eq!(
a_query_with_b.iter_combinations::<2>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_with_b.iter_combinations::<3>(&world).count(), 0);
assert_eq!(
a_query_with_b.iter_combinations::<3>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_with_b.iter_combinations::<4>(&world).count(), 0);
assert_eq!(
a_query_with_b.iter_combinations::<4>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_with_b.iter_combinations::<5>(&world).count(), 0);
assert_eq!(
a_query_with_b.iter_combinations::<5>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_with_b.iter_combinations::<1024>(&world).count(), 0);
assert_eq!(
a_query_with_b.iter_combinations::<1024>(&world).size_hint(),
(0, Some(0))
);
let mut a_with_b = world.query_filtered::<&A, With<B>>();
let w = &world;
assert_eq!(a_with_b.iter_combinations::<0>(w).count(), 0);
assert_eq!(a_with_b.iter_combinations::<0>(w).size_hint().1, Some(0));
assert_eq!(a_with_b.iter_combinations::<1>(w).count(), 1);
assert_eq!(a_with_b.iter_combinations::<1>(w).size_hint().1, Some(1));
assert_eq!(a_with_b.iter_combinations::<2>(w).count(), 0);
assert_eq!(a_with_b.iter_combinations::<2>(w).size_hint().1, Some(0));
assert_eq!(a_with_b.iter_combinations::<3>(w).count(), 0);
assert_eq!(a_with_b.iter_combinations::<3>(w).size_hint().1, Some(0));
assert_eq!(a_with_b.iter_combinations::<4>(w).count(), 0);
assert_eq!(a_with_b.iter_combinations::<4>(w).size_hint().1, Some(0));
assert_eq!(a_with_b.iter_combinations::<5>(w).count(), 0);
assert_eq!(a_with_b.iter_combinations::<5>(w).size_hint().1, Some(0));
assert_eq!(a_with_b.iter_combinations::<1024>(w).count(), 0);
assert_eq!(a_with_b.iter_combinations::<1024>(w).size_hint().1, Some(0));
let mut a_query_without_b = world.query_filtered::<&A, Without<B>>();
assert_eq!(a_query_without_b.iter_combinations::<0>(&world).count(), 0);
assert_eq!(
a_query_without_b.iter_combinations::<0>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_without_b.iter_combinations::<1>(&world).count(), 3);
assert_eq!(
a_query_without_b.iter_combinations::<1>(&world).size_hint(),
(0, Some(3))
);
assert_eq!(a_query_without_b.iter_combinations::<2>(&world).count(), 3);
assert_eq!(
a_query_without_b.iter_combinations::<2>(&world).size_hint(),
(0, Some(3))
);
assert_eq!(a_query_without_b.iter_combinations::<3>(&world).count(), 1);
assert_eq!(
a_query_without_b.iter_combinations::<3>(&world).size_hint(),
(0, Some(1))
);
assert_eq!(a_query_without_b.iter_combinations::<4>(&world).count(), 0);
assert_eq!(
a_query_without_b.iter_combinations::<4>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query_without_b.iter_combinations::<5>(&world).count(), 0);
assert_eq!(
a_query_without_b.iter_combinations::<5>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(
a_query_without_b.iter_combinations::<1024>(&world).count(),
0
);
assert_eq!(
a_query_without_b
.iter_combinations::<1024>(&world)
.size_hint(),
(0, Some(0))
);
let mut a_wout_b = world.query_filtered::<&A, Without<B>>();
let w = &world;
assert_eq!(a_wout_b.iter_combinations::<0>(w).count(), 0);
assert_eq!(a_wout_b.iter_combinations::<0>(w).size_hint().1, Some(0));
assert_eq!(a_wout_b.iter_combinations::<1>(w).count(), 3);
assert_eq!(a_wout_b.iter_combinations::<1>(w).size_hint().1, Some(3));
assert_eq!(a_wout_b.iter_combinations::<2>(w).count(), 3);
assert_eq!(a_wout_b.iter_combinations::<2>(w).size_hint().1, Some(3));
assert_eq!(a_wout_b.iter_combinations::<3>(w).count(), 1);
assert_eq!(a_wout_b.iter_combinations::<3>(w).size_hint().1, Some(1));
assert_eq!(a_wout_b.iter_combinations::<4>(w).count(), 0);
assert_eq!(a_wout_b.iter_combinations::<4>(w).size_hint().1, Some(0));
assert_eq!(a_wout_b.iter_combinations::<5>(w).count(), 0);
assert_eq!(a_wout_b.iter_combinations::<5>(w).size_hint().1, Some(0));
assert_eq!(a_wout_b.iter_combinations::<1024>(w).count(), 0);
assert_eq!(a_wout_b.iter_combinations::<1024>(w).size_hint().1, Some(0));
let values: HashSet<[&A; 2]> = a_query_without_b.iter_combinations(&world).collect();
let values: HashSet<[&A; 2]> = a_wout_b.iter_combinations(&world).collect();
assert_eq!(
values,
[[&A(2), &A(3)], [&A(2), &A(4)], [&A(3), &A(4)],]
@ -239,7 +174,7 @@ mod tests {
.collect::<HashSet<_>>()
);
let values: HashSet<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
let values: HashSet<[&A; 3]> = a_wout_b.iter_combinations(&world).collect();
assert_eq!(
values,
[[&A(2), &A(3), &A(4)],].into_iter().collect::<HashSet<_>>()
@ -269,7 +204,7 @@ mod tests {
c.0 += 1000;
}
let values: HashSet<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
let values: HashSet<[&A; 3]> = a_wout_b.iter_combinations(&world).collect();
assert_eq!(
values,
[[&A(12), &A(103), &A(1004)],]