[Merged by Bors] - Allow iter combinations on queries with filters

crates/bevy_ecs/src/query/filter.rs

@@ -166,6 +166,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`].
@@ -289,6 +299,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
@@ -319,9 +339,11 @@ 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(Clone, Copy)]
 pub struct OrFetch<T: FilterFetch> {
     fetch: T,
     matches: bool,
@@ -586,6 +608,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> {}
     };
 }
 

crates/bevy_ecs/src/query/mod.rs

@@ -12,9 +12,24 @@ pub use state::*;
 
 #[cfg(test)]
 mod tests {
+    macro_rules! hash_set {
+        ( $($key:expr,)+) => { hash_set!($($key),+) };
+        ( $($key:expr),* ) => {
+            {
+                let mut _set = ::std::collections::HashSet::new();
+                $(
+                    let _ = _set.insert($key);
+                )*
+                _set
+            }
+        }
+    }
+
+    use std::collections::HashSet;
+
     use crate::{self as bevy_ecs, component::Component, world::World};
 
-    #[derive(Component, Debug, Eq, PartialEq)]
+    #[derive(Component, Debug, Hash, Eq, PartialEq)]
     struct A(usize);
     #[derive(Component, Debug, Eq, PartialEq)]
     struct B(usize);
@@ -137,6 +152,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: HashSet<[&A; 2]> = a_query_without_b.iter_combinations(&world).collect();
+        assert_eq!(
+            values,
+            hash_set![[&A(2), &A(3)], [&A(2), &A(4)], [&A(3), &A(4)],]
+        );
+
+        let values: HashSet<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
+        assert_eq!(values, hash_set![[&A(2), &A(3), &A(4)],]);
+
+        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,
+            hash_set![
+                [&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: HashSet<[&A; 3]> = a_query_without_b.iter_combinations(&world).collect();
+        assert_eq!(values, hash_set![[&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: HashSet<[&A; 3]> = query_changed.iter_combinations(&world).collect();
+        assert_eq!(
+            values,
+            hash_set![
+                [&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();