diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index 1626dd43fc0138..ad79241b109800 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -739,3 +739,30 @@ impl_tick_filter!( ChangedFetch, ComponentTicks::is_changed ); + +/// A marker trait to indicate that the filter works at an archetype level. +/// +/// This is needed to implement [`ExactSizeIterator`](std::iter::ExactSizeIterator) for +/// [`QueryIter`](crate::query::QueryIter) that contains archetype-level filters. +/// +/// The trait must only be implement for filters where its corresponding [`Fetch::IS_ARCHETYPAL`](crate::query::Fetch::IS_ARCHETYPAL) +/// is [`prim@true`]. As such, only the [`With`] and [`Without`] filters can implement the trait. +/// [Tuples](prim@tuple) and [`Or`] filters are automatically implemented with the trait only if its containing types +/// also implement the same trait. +/// +/// [`Added`] and [`Changed`] works with entities, and therefore are not archetypal. As such +/// they do not implement [`ArchetypeFilter`]. +pub trait ArchetypeFilter {} + +impl ArchetypeFilter for With {} +impl ArchetypeFilter for Without {} + +macro_rules! impl_archetype_filter_tuple { + ($($filter: ident),*) => { + impl<$($filter: ArchetypeFilter),*> ArchetypeFilter for ($($filter,)*) {} + + impl<$($filter: ArchetypeFilter),*> ArchetypeFilter for Or<($($filter,)*)> {} + }; +} + +all_tuples!(impl_archetype_filter_tuple, 0, 15, F); diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index b3c304f231aa6c..58c6d40b2a7086 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -2,7 +2,7 @@ use crate::{ archetype::{ArchetypeId, Archetypes}, entity::{Entities, Entity}, prelude::World, - query::{Fetch, QueryState, WorldQuery}, + query::{ArchetypeFilter, Fetch, QueryState, WorldQuery}, storage::{TableId, Tables}, }; use std::{borrow::Borrow, iter::FusedIterator, marker::PhantomData, mem::MaybeUninit}; @@ -346,15 +346,10 @@ where } } -// NOTE: We can cheaply implement this for unfiltered Queries because we have: -// (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. -// 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, ()> +impl<'w, 's, Q: WorldQuery, QF, F> ExactSizeIterator for QueryIter<'w, 's, Q, QF, F> where QF: Fetch<'w, State = Q::State>, + F: WorldQuery + ArchetypeFilter, { fn len(&self) -> usize { self.query_state diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index d037c4a5419745..74de0fd802bee7 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -31,6 +31,8 @@ mod tests { struct B(usize); #[derive(Component, Debug, Eq, PartialEq, Clone, Copy)] struct C(usize); + #[derive(Component, Debug, Eq, PartialEq, Clone, Copy)] + struct D(usize); #[derive(Component, Debug, Eq, PartialEq, Clone, Copy)] #[component(storage = "SparseSet")] @@ -51,6 +53,145 @@ mod tests { assert_eq!(values, vec![&B(3)]); } + #[test] + fn query_filtered_len() { + let mut world = World::new(); + world.spawn().insert_bundle((A(1), B(1))); + world.spawn().insert_bundle((A(2),)); + world.spawn().insert_bundle((A(3),)); + + let mut values = world.query_filtered::<&A, With>(); + let n = 1; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Without>(); + let n = 2; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + + let mut world = World::new(); + world.spawn().insert_bundle((A(1), B(1), C(1))); + world.spawn().insert_bundle((A(2), B(2))); + world.spawn().insert_bundle((A(3), B(3))); + world.spawn().insert_bundle((A(4), C(4))); + world.spawn().insert_bundle((A(5), C(5))); + world.spawn().insert_bundle((A(6), C(6))); + world.spawn().insert_bundle((A(7),)); + world.spawn().insert_bundle((A(8),)); + world.spawn().insert_bundle((A(9),)); + world.spawn().insert_bundle((A(10),)); + + // With/Without for B and C + let mut values = world.query_filtered::<&A, With>(); + let n = 3; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, With>(); + let n = 4; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Without>(); + let n = 7; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Without>(); + let n = 6; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + + // With/Without (And) combinations + let mut values = world.query_filtered::<&A, (With, With)>(); + let n = 1; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, (With, Without)>(); + let n = 2; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, (Without, With)>(); + let n = 3; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, (Without, Without)>(); + let n = 4; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + + // With/Without Or<()> combinations + let mut values = world.query_filtered::<&A, Or<(With, With)>>(); + let n = 6; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Or<(With, Without)>>(); + let n = 7; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Or<(Without, With)>>(); + let n = 8; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Or<(Without, Without)>>(); + let n = 9; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + + let mut values = world.query_filtered::<&A, (Or<(With,)>, Or<(With,)>)>(); + let n = 1; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Or<(Or<(With, With)>, With)>>(); + let n = 6; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + + world.spawn().insert_bundle((A(11), D(11))); + + let mut values = world.query_filtered::<&A, Or<(Or<(With, With)>, With)>>(); + let n = 7; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + let mut values = world.query_filtered::<&A, Or<(Or<(With, With)>, Without)>>(); + let n = 10; + assert_eq!(values.iter(&world).size_hint().0, n); + assert_eq!(values.iter(&world).size_hint().1.unwrap(), n); + assert_eq!(values.iter(&world).len(), n); + assert_eq!(values.iter(&world).count(), n); + } + #[test] fn query_iter_combinations() { let mut world = World::new(); diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.rs new file mode 100644 index 00000000000000..e634bd86d217c6 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.rs @@ -0,0 +1,18 @@ +use bevy_ecs::prelude::*; + +#[derive(Component)] +struct Foo; + +fn on_changed(query: Query<&Foo, Changed>) { + // this should fail to compile + is_exact_size_iterator(query.iter()); +} + +fn on_added(query: Query<&Foo, Added>) { + // this should fail to compile + is_exact_size_iterator(query.iter()); +} + +fn is_exact_size_iterator(_iter: T) {} + +fn main() {} \ No newline at end of file diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.stderr new file mode 100644 index 00000000000000..22ea39040d131f --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.stderr @@ -0,0 +1,29 @@ +error[E0277]: the trait bound `bevy_ecs::query::Changed: ArchetypeFilter` is not satisfied + --> tests/ui/query_exact_sized_iterator_safety.rs:8:28 + | +8 | is_exact_size_iterator(query.iter()); + | ---------------------- ^^^^^^^^^^^^ the trait `ArchetypeFilter` is not implemented for `bevy_ecs::query::Changed` + | | + | required by a bound introduced by this call + | + = note: required because of the requirements on the impl of `ExactSizeIterator` for `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Changed>` +note: required by a bound in `is_exact_size_iterator` + --> tests/ui/query_exact_sized_iterator_safety.rs:16:30 + | +16 | fn is_exact_size_iterator(_iter: T) {} + | ^^^^^^^^^^^^^^^^^ required by this bound in `is_exact_size_iterator` + +error[E0277]: the trait bound `bevy_ecs::query::Added: ArchetypeFilter` is not satisfied + --> tests/ui/query_exact_sized_iterator_safety.rs:13:28 + | +13 | is_exact_size_iterator(query.iter()); + | ---------------------- ^^^^^^^^^^^^ the trait `ArchetypeFilter` is not implemented for `bevy_ecs::query::Added` + | | + | required by a bound introduced by this call + | + = note: required because of the requirements on the impl of `ExactSizeIterator` for `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Added>` +note: required by a bound in `is_exact_size_iterator` + --> tests/ui/query_exact_sized_iterator_safety.rs:16:30 + | +16 | fn is_exact_size_iterator(_iter: T) {} + | ^^^^^^^^^^^^^^^^^ required by this bound in `is_exact_size_iterator`