Update ExactSizeIterator impl to support archetypal filters (With, …

…Without) (#5124)

# Objective

- Fixes #3142

## Solution

- Done according to #3142
- Created new marker trait `ArchetypeFilter`
- Implement said trait to:
  - `With<T>`
  - `Without<T>`
  - tuples containing only types that implement `ArchetypeFilter`, from 0 to 15 elements
  - `Or<T>` where T is a tuple as described previously
- Changed `ExactSizeIterator` impl to include a new generic that must implement `WorldQuery` and `ArchetypeFilter`
- Added new tests

---

## Changelog

### Added
- `Query`s with archetypal filters can now use `.iter().len()` to get the exact size of the iterator.

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

@@ -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<B>>();
+        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<B>>();
+        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<B>>();
+        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<C>>();
+        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<B>>();
+        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<C>>();
+        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<B>, With<C>)>();
+        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<B>, Without<C>)>();
+        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<B>, With<C>)>();
+        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<B>, Without<C>)>();
+        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<B>, With<C>)>>();
+        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<B>, Without<C>)>>();
+        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<B>, With<C>)>>();
+        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<B>, Without<C>)>>();
+        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<B>,)>, Or<(With<C>,)>)>();
+        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<B>, With<C>)>, With<D>)>>();
+        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<B>, With<C>)>, With<D>)>>();
+        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<B>, With<C>)>, Without<D>)>>();
+        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();

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<Foo>>) {
+    // this should fail to compile
+    is_exact_size_iterator(query.iter());
+}
+
+fn on_added(query: Query<&Foo, Added<Foo>>) {
+    // this should fail to compile
+    is_exact_size_iterator(query.iter());
+}
+
+fn is_exact_size_iterator<T: ExactSizeIterator>(_iter: T) {}
+
+fn main() {}

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<Foo>: 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<Foo>`
+   |     |
+   |     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<Foo>>`
+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<T: ExactSizeIterator>(_iter: T) {}
+   |                              ^^^^^^^^^^^^^^^^^ required by this bound in `is_exact_size_iterator`
+
+error[E0277]: the trait bound `bevy_ecs::query::Added<Foo>: 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<Foo>`
+   |     |
+   |     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<Foo>>`
+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<T: ExactSizeIterator>(_iter: T) {}
+   |                              ^^^^^^^^^^^^^^^^^ required by this bound in `is_exact_size_iterator`