Skip to content

Commit

Permalink
Update ExactSizeIterator impl to support archetypal filters (With, …
Browse files Browse the repository at this point in the history
…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.
  • Loading branch information
harudagondi committed Jun 29, 2022
1 parent e28b88b commit 6e50b24
Show file tree
Hide file tree
Showing 5 changed files with 218 additions and 8 deletions.
27 changes: 27 additions & 0 deletions crates/bevy_ecs/src/query/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
11 changes: 3 additions & 8 deletions crates/bevy_ecs/src/query/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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
Expand Down
141 changes: 141 additions & 0 deletions crates/bevy_ecs/src/query/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand All @@ -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();
Expand Down
Original file line number Diff line number Diff line change
@@ -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() {}
Original file line number Diff line number Diff line change
@@ -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`

0 comments on commit 6e50b24

Please sign in to comment.