Skip to content

Commit

Permalink
Fix iter_combinations with filters (#3651)
Browse files Browse the repository at this point in the history
  • Loading branch information
harudagondi committed Jan 13, 2022
1 parent 84144c9 commit bf7b9ca
Show file tree
Hide file tree
Showing 2 changed files with 220 additions and 0 deletions.
44 changes: 44 additions & 0 deletions crates/bevy_ecs/src/query/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ impl<T: Component> WorldQuery for With<T> {
}

/// The [`Fetch`] of [`With`].
#[derive(Copy)]
pub struct WithFetch<T> {
marker: PhantomData<T>,
}
Expand Down Expand Up @@ -166,6 +167,14 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch<T> {
// SAFETY: no component access or archetype component access
unsafe impl<T> ReadOnlyFetch for WithFetch<T> {}

impl<T> Clone for WithFetch<T> {
fn clone(&self) -> Self {
Self {
marker: self.marker,
}
}
}

/// Filter that selects entities without a component `T`.
///
/// This is the negation of [`With`].
Expand Down Expand Up @@ -199,6 +208,7 @@ impl<T: Component> WorldQuery for Without<T> {
}

/// The [`Fetch`] of [`Without`].
#[derive(Copy)]
pub struct WithoutFetch<T> {
marker: PhantomData<T>,
}
Expand Down Expand Up @@ -289,6 +299,14 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithoutFetch<T> {
// SAFETY: no component access or archetype component access
unsafe impl<T> ReadOnlyFetch for WithoutFetch<T> {}

impl<T> Clone for WithoutFetch<T> {
fn clone(&self) -> Self {
Self {
marker: self.marker,
}
}
}

/// A filter that tests if any of the given filters apply.
///
/// This is useful for example if a system with multiple components in a query only wants to run
Expand Down Expand Up @@ -319,14 +337,25 @@ unsafe impl<T> ReadOnlyFetch for WithoutFetch<T> {}
/// }
/// # print_cool_entity_system.system();
/// ```
#[derive(Clone, Copy)]
pub struct Or<T>(pub T);

/// The [`Fetch`] of [`Or`].
#[derive(Copy)]
pub struct OrFetch<T: FilterFetch> {
fetch: T,
matches: bool,
}

impl<T: FilterFetch + Clone> Clone for OrFetch<T> {
fn clone(&self) -> Self {
Self {
fetch: self.fetch.clone(),
matches: self.matches,
}
}
}

macro_rules! impl_query_filter_tuple {
($(($filter: ident, $state: ident)),*) => {
#[allow(unused_variables)]
Expand Down Expand Up @@ -456,6 +485,7 @@ macro_rules! impl_tick_filter {
pub struct $name<T>(PhantomData<T>);

$(#[$fetch_meta])*
#[derive(Copy)]
pub struct $fetch_name<T> {
table_ticks: *const UnsafeCell<ComponentTicks>,
entity_table_rows: *const usize,
Expand Down Expand Up @@ -586,6 +616,20 @@ macro_rules! impl_tick_filter {

/// SAFETY: read-only access
unsafe impl<T: Component> ReadOnlyFetch for $fetch_name<T> {}

impl<T> Clone for $fetch_name<T> {
fn clone(&self) -> Self {
Self {
table_ticks: self.table_ticks.clone(),
entity_table_rows: self.entity_table_rows.clone(),
marker: self.marker.clone(),
entities: self.entities.clone(),
sparse_set: self.sparse_set.clone(),
last_change_tick: self.last_change_tick.clone(),
change_tick: self.change_tick.clone(),
}
}
}
};
}

Expand Down
176 changes: 176 additions & 0 deletions crates/bevy_ecs/src/query/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,182 @@ mod tests {
assert_eq!(values, Vec::<[&B; 2]>::new());
}

#[test]
fn query_filtered_iter_combinations() {
use bevy_ecs::query::{Added, Changed, Or, With, Without};

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),));
world.spawn().insert_bundle((A(4),));

let mut a_query_with_b = world.query_filtered::<&A, With<B>>();
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_query_without_b = world.query_filtered::<&A, Without<B>>();
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 values: Vec<[&A; 2]> = a_query_without_b.iter_combinations(&world).collect();
assert_eq!(
values,
vec![[&A(2), &A(3)], [&A(2), &A(4)], [&A(3), &A(4)],]
);

let values: Vec<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
assert_eq!(values, vec![[&A(2), &A(3), &A(4)],]);

let mut query = world.query_filtered::<&A, Or<(With<A>, With<B>)>>();
let values: Vec<[&A; 2]> = query.iter_combinations(&world).collect();
assert_eq!(
values,
vec![
[&A(1), &A(2)],
[&A(1), &A(3)],
[&A(1), &A(4)],
[&A(2), &A(3)],
[&A(2), &A(4)],
[&A(3), &A(4)],
]
);

let mut query = world.query_filtered::<&mut A, Without<B>>();
let mut combinations = query.iter_combinations_mut(&mut world);
while let Some([mut a, mut b, mut c]) = combinations.fetch_next() {
a.0 += 10;
b.0 += 100;
c.0 += 1000;
}

let values: Vec<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
assert_eq!(values, vec![[&A(12), &A(103), &A(1004)],]);

// Check if Added<T>, Changed<T> works
let mut world = World::new();

world.spawn().insert_bundle((A(1), B(1)));
world.spawn().insert_bundle((A(2), B(2)));
world.spawn().insert_bundle((A(3), B(3)));
world.spawn().insert_bundle((A(4), B(4)));

let mut query_added = world.query_filtered::<&A, Added<A>>();

world.clear_trackers();
world.spawn().insert_bundle((A(5),));

assert_eq!(query_added.iter_combinations::<2>(&world).count(), 0);

world.clear_trackers();
world.spawn().insert_bundle((A(6),));
world.spawn().insert_bundle((A(7),));

assert_eq!(query_added.iter_combinations::<2>(&world).count(), 1);

world.clear_trackers();
world.spawn().insert_bundle((A(8),));
world.spawn().insert_bundle((A(9),));
world.spawn().insert_bundle((A(10),));

assert_eq!(query_added.iter_combinations::<2>(&world).count(), 3);

world.clear_trackers();

let mut query_changed = world.query_filtered::<&A, Changed<A>>();

let mut query = world.query_filtered::<&mut A, With<B>>();
let mut combinations = query.iter_combinations_mut(&mut world);
while let Some([mut a, mut b, mut c]) = combinations.fetch_next() {
a.0 += 10;
b.0 += 100;
c.0 += 1000;
}

let values: Vec<[&A; 3]> = query_changed.iter_combinations(&world).collect();
assert_eq!(
values,
vec![
[&A(31), &A(212), &A(1203)],
[&A(31), &A(212), &A(3004)],
[&A(31), &A(1203), &A(3004)],
[&A(212), &A(1203), &A(3004)]
]
);
}

#[test]
fn query_iter_combinations_sparse() {
let mut world = World::new();
Expand Down

0 comments on commit bf7b9ca

Please sign in to comment.