[Merged by Bors] - Update ExactSizeIterator impl to support archetypal filters (With, Without)

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<T> ArchetypeFilter for With<T> {}
+impl<T> ArchetypeFilter for Without<T> {}
+
+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);

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<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, ()>
+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

crates/bevy_ecs/src/query/mod.rs

@@ -51,6 +51,72 @@ 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),));
+
+        fn are_sizes_equal_to(iter: impl ExactSizeIterator, n: usize) -> bool {
+            [
+                iter.size_hint().0,
+                iter.size_hint().1.unwrap(),
+                iter.len(),
+                iter.count(),
+            ]
+            .into_iter()
+            .all(|x| x == n)
+        }
+
+        let mut values = world.query_filtered::<&A, With<B>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 1));
+        let mut values = world.query_filtered::<&A, Without<B>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 2));
+
+        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<B>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 3));
+        let mut values = world.query_filtered::<&A, With<C>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 4));
+        let mut values = world.query_filtered::<&A, Without<B>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 7));
+        let mut values = world.query_filtered::<&A, Without<C>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 6));
+
+        // With/Without (And) combinations
+        let mut values = world.query_filtered::<&A, (With<B>, With<C>)>();
+        assert!(are_sizes_equal_to(values.iter(&world), 1));
+        let mut values = world.query_filtered::<&A, (With<B>, Without<C>)>();
+        assert!(are_sizes_equal_to(values.iter(&world), 2));
+        let mut values = world.query_filtered::<&A, (Without<B>, With<C>)>();
+        assert!(are_sizes_equal_to(values.iter(&world), 3));
+        let mut values = world.query_filtered::<&A, (Without<B>, Without<C>)>();
+        assert!(are_sizes_equal_to(values.iter(&world), 4));
+
+        // With/Without Or<()> combinations
+        let mut values = world.query_filtered::<&A, Or<(With<B>, With<C>)>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 6));
+        let mut values = world.query_filtered::<&A, Or<(With<B>, Without<C>)>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 7));
+        let mut values = world.query_filtered::<&A, Or<(Without<B>, With<C>)>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 8));
+        let mut values = world.query_filtered::<&A, Or<(Without<B>, Without<C>)>>();
+        assert!(are_sizes_equal_to(values.iter(&world), 9));
+    }
+
     #[test]
     fn query_iter_combinations() {
         let mut world = World::new();

crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.rs

@@ -0,0 +1,16 @@
+use bevy_ecs::prelude::*;
+
+#[derive(Component)]
+struct Foo;
+
+fn on_changed(query: Query<&Foo, Changed<Foo>>) {
+    // this should fail to compile
+    println!("{}", query.iter().len())
+}
+
+fn on_added(query: Query<&Foo, Added<Foo>>) {
+    // this should fail to compile
+    println!("{}", query.iter().len())
+}
+
+fn main() {}

crates/bevy_ecs_compile_fail_tests/tests/ui/query_exact_sized_iterator_safety.stderr

@@ -0,0 +1,50 @@
+error[E0599]: the method `len` exists for struct `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Changed<Foo>>`, but its trait bounds were not satisfied
+   --> tests/ui/query_exact_sized_iterator_safety.rs:8:33
+    |
+8   |       println!("{}", query.iter().len())
+    |                                   ^^^ method cannot be called on `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Changed<Foo>>` due to unsatisfied trait bounds
+    |
+   ::: /home/runner/work/bevy/bevy/crates/bevy_ecs/src/query/iter.rs:16:1
+    |
+16  |   pub struct QueryIter<'w, 's, Q: WorldQuery, QF: Fetch<'w, State = Q::State>, F: WorldQuery> {
+    |   ------------------------------------------------------------------------------------------- doesn't satisfy `_: ExactSizeIterator`
+    |
+   ::: /home/runner/work/bevy/bevy/crates/bevy_ecs/src/query/filter.rs:703:1
+    |
+703 | / impl_tick_filter!(
+704 | |     /// A filter on a component that only retains results added or mutably dereferenced after the system last ran.
+705 | |     ///
+706 | |     /// A common use for this filter is avoiding redundant work when values have not changed.
+...   |
+740 | |     ComponentTicks::is_changed
+741 | | );
+    | |_- doesn't satisfy `bevy_ecs::query::Changed<Foo>: ArchetypeFilter`
+    |
+    = note: the following trait bounds were not satisfied:
+            `bevy_ecs::query::Changed<Foo>: ArchetypeFilter`
+            which is required by `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Changed<Foo>>: ExactSizeIterator`
+error[E0599]: the method `len` exists for struct `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Added<Foo>>`, but its trait bounds were not satisfied
+   --> tests/ui/query_exact_sized_iterator_safety.rs:13:33
+    |
+13  |       println!("{}", query.iter().len())
+    |                                   ^^^ method cannot be called on `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Added<Foo>>` due to unsatisfied trait bounds
+    |
+   ::: /home/runner/work/bevy/bevy/crates/bevy_ecs/src/query/iter.rs:16:1
+    |
+16  |   pub struct QueryIter<'w, 's, Q: WorldQuery, QF: Fetch<'w, State = Q::State>, F: WorldQuery> {
+    |   ------------------------------------------------------------------------------------------- doesn't satisfy `_: ExactSizeIterator`
+    |
+   ::: /home/runner/work/bevy/bevy/crates/bevy_ecs/src/query/filter.rs:668:1
+    |
+668 | / impl_tick_filter!(
+669 | |     /// A filter on a component that only retains results added after the system last ran.
+670 | |     ///
+671 | |     /// A common use for this filter is one-time initialization.
+...   |
+700 | |     ComponentTicks::is_added
+701 | | );
+    | |_- doesn't satisfy `bevy_ecs::query::Added<Foo>: ArchetypeFilter`
+    |
+    = note: the following trait bounds were not satisfied:
+            `bevy_ecs::query::Added<Foo>: ArchetypeFilter`
+            which is required by `QueryIter<'_, '_, &Foo, ReadFetch<'_, Foo>, bevy_ecs::query::Added<Foo>>: ExactSizeIterator`