Skip to content

Commit

Permalink
Allow iter combinations on queries with filters (bevyengine#3656)
Browse files Browse the repository at this point in the history 
# Objective

- Previously, `iter_combinations()` does not work on queries that have filters.
- Fixes bevyengine#3651

## Solution

- Derived Copy on all `*Fetch<T>` structs, and manually implemented `Clone` to allow the test to pass (`.count()` does not work on `QueryCombinationIter` when `Clone` is derived)


Co-authored-by: Carter Anderson <[email protected]>
  • Loading branch information
@harudagondi @cart @aevyrie
2 people authored and aevyrie committed Jun 7, 2022
1 parent f1a6f4d commit c9f7a47
Show file tree
Hide file tree
Showing 2 changed files with 230 additions and 1 deletion.
38 changes: 38 additions & 0 deletions crates/bevy_ecs/src/query/filter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,16 @@ 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,
}
}
}

impl<T> Copy for WithFetch<T> {}

/// Filter that selects entities without a component `T`.
///
/// This is the negation of [`With`].
Expand Down Expand Up @@ -296,6 +306,16 @@ 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,
}
}
}

impl<T> Copy for WithoutFetch<T> {}

/// 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 @@ -326,9 +346,11 @@ unsafe impl<T> ReadOnlyFetch for WithoutFetch<T> {}
/// }
/// # bevy_ecs::system::assert_is_system(print_cool_entity_system);
/// ```
#[derive(Clone, Copy)]
pub struct Or<T>(pub T);

/// The [`Fetch`] of [`Or`].
#[derive(Clone, Copy)]
#[doc(hidden)]
pub struct OrFetch<T: FilterFetch> {
fetch: T,
Expand Down Expand Up @@ -596,6 +618,22 @@ 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(),
}
}
}

impl<T> Copy for $fetch_name<T> {}
};
}

Expand Down
193 changes: 192 additions & 1 deletion crates/bevy_ecs/src/query/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,9 @@ pub use state::*;
mod tests {
use super::AnyOf;
use crate::{self as bevy_ecs, component::Component, world::World};
use std::collections::HashSet;

#[derive(Component, Debug, Eq, PartialEq)]
#[derive(Component, Debug, Hash, Eq, PartialEq)]
struct A(usize);
#[derive(Component, Debug, Eq, PartialEq)]
struct B(usize);
Expand Down Expand Up @@ -140,6 +141,196 @@ 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: HashSet<[&A; 2]> = a_query_without_b.iter_combinations(&world).collect();
assert_eq!(
values,
[[&A(2), &A(3)], [&A(2), &A(4)], [&A(3), &A(4)],]
.into_iter()
.collect::<HashSet<_>>()
);

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

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

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: HashSet<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
assert_eq!(
values,
[[&A(12), &A(103), &A(1004)],]
.into_iter()
.collect::<HashSet<_>>()
);

// 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: HashSet<[&A; 3]> = query_changed.iter_combinations(&world).collect();
assert_eq!(
values,
[
[&A(31), &A(212), &A(1203)],
[&A(31), &A(212), &A(3004)],
[&A(31), &A(1203), &A(3004)],
[&A(212), &A(1203), &A(3004)]
]
.into_iter()
.collect::<HashSet<_>>()
);
}

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

0 comments on commit c9f7a47

Please sign in to comment.