From c9f7a4732ac1312cdddc0e98f3bfa4257eb0f592 Mon Sep 17 00:00:00 2001 From: harudagondi Date: Fri, 8 Apr 2022 00:21:24 +0000 Subject: [PATCH] Allow iter combinations on queries with filters (#3656) # Objective - Previously, `iter_combinations()` does not work on queries that have filters. - Fixes #3651 ## Solution - Derived Copy on all `*Fetch` 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 --- crates/bevy_ecs/src/query/filter.rs | 38 ++++++ crates/bevy_ecs/src/query/mod.rs | 193 +++++++++++++++++++++++++++- 2 files changed, 230 insertions(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index dd12aec8ef8429..0cf2e79ff80879 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -171,6 +171,16 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch { // SAFETY: no component access or archetype component access unsafe impl ReadOnlyFetch for WithFetch {} +impl Clone for WithFetch { + fn clone(&self) -> Self { + Self { + marker: self.marker, + } + } +} + +impl Copy for WithFetch {} + /// Filter that selects entities without a component `T`. /// /// This is the negation of [`With`]. @@ -296,6 +306,16 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithoutFetch { // SAFETY: no component access or archetype component access unsafe impl ReadOnlyFetch for WithoutFetch {} +impl Clone for WithoutFetch { + fn clone(&self) -> Self { + Self { + marker: self.marker, + } + } +} + +impl Copy for WithoutFetch {} + /// 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 @@ -326,9 +346,11 @@ unsafe impl ReadOnlyFetch for WithoutFetch {} /// } /// # bevy_ecs::system::assert_is_system(print_cool_entity_system); /// ``` +#[derive(Clone, Copy)] pub struct Or(pub T); /// The [`Fetch`] of [`Or`]. +#[derive(Clone, Copy)] #[doc(hidden)] pub struct OrFetch { fetch: T, @@ -596,6 +618,22 @@ macro_rules! impl_tick_filter { /// SAFETY: read-only access unsafe impl ReadOnlyFetch for $fetch_name {} + + impl Clone for $fetch_name { + 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 Copy for $fetch_name {} }; } diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 0603a53bfd2ae8..191b3dd3497a78 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -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); @@ -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>(); + 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>(); + 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::>() + ); + + 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::>() + ); + + let mut query = world.query_filtered::<&A, Or<(With, With)>>(); + 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::>() + ); + + let mut query = world.query_filtered::<&mut A, Without>(); + 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::>() + ); + + // Check if Added, Changed 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>(); + + 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>(); + + let mut query = world.query_filtered::<&mut A, With>(); + 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::>() + ); + } + #[test] fn query_iter_combinations_sparse() { let mut world = World::new();