From 71a246ce9e4d25c98e3cadc246e0735e98fe8f41 Mon Sep 17 00:00:00 2001 From: Nicola Papale Date: Wed, 27 Apr 2022 18:02:06 +0000 Subject: [PATCH] Improve QueryIter size_hint hints (#4244) ## Objective This fixes #1686. `size_hint` can be useful even if a little niche. For example, `collect::>()` 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. --- crates/bevy_ecs/macros/src/fetch.rs | 4 + crates/bevy_ecs/src/lib.rs | 37 ++++++- crates/bevy_ecs/src/query/fetch.rs | 25 +++++ crates/bevy_ecs/src/query/filter.rs | 8 ++ crates/bevy_ecs/src/query/iter.rs | 17 ++- crates/bevy_ecs/src/query/mod.rs | 165 +++++++++------------------- 6 files changed, 131 insertions(+), 125 deletions(-) diff --git a/crates/bevy_ecs/macros/src/fetch.rs b/crates/bevy_ecs/macros/src/fetch.rs index 3b13a8c609..4b07fc2a4e 100644 --- a/crates/bevy_ecs/macros/src/fetch.rs +++ b/crates/bevy_ecs/macros/src/fetch.rs @@ -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` diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 857af11972..50edc31427 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -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, Without)],); + assert_eq!(3, query_min_size![&B, Or<(With, With)>],); + assert_eq!(1, query_min_size![&B, (With, With)],); + assert_eq!(1, query_min_size![(&A, &B), With],); + assert_eq!(4, query_min_size![&A, ()], "Simple Archetypal"); + assert_eq!(4, query_min_size![ChangeTrackers, ()],); + // 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], "Simple Added"); + assert_eq!(0, query_min_size![(), Changed], "Simple Changed"); + assert_eq!(0, query_min_size![(&A, &B), Changed],); + assert_eq!(0, query_min_size![&A, (Changed, With)],); + assert_eq!(0, query_min_size![(&A, &B), Or<(Changed, Changed)>],); + } + #[test] fn reserve_entities_across_worlds() { let mut world_a = World::default(); diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index d68c64375e..5913757f05 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -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 { } }; + 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 { } }; + 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 { } }; + 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 { 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 { } }; + 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 { const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; + #[inline(always)] unsafe fn init( _world: &World, diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index 0cf2e79ff8..d1218f5406 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -145,6 +145,8 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch { } }; + 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 { } }; + 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() diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 7e135c7808..71704b7f61 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -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) { 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) { 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 +// TODO: add an ArchetypeOnlyFilter that enables us to implement this for filters like With. +// 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>, diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 191b3dd349..05f87a73bc 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -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>(); - 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>(); + 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>(); - 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>(); + 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::>() ); - 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::>() @@ -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)],]