Add Joinable trait allowing filtering the results of one query with another

crates/bevy_ecs/src/query/iter.rs

@@ -1,12 +1,14 @@
 use crate::{
     archetype::{ArchetypeId, Archetypes},
-    query::{Fetch, QueryState, WorldQuery},
+    entity::Entity,
+    query::{Fetch, QueryState, ReadOnlyFetch, WorldQuery},
     storage::{TableId, Tables},
+    system::Query,
     world::World,
 };
 use std::{marker::PhantomData, mem::MaybeUninit};
 
-use super::{QueryFetch, QueryItem, ReadOnlyFetch};
+use super::{QueryFetch, QueryItem, ROQueryItem, WorldQueryGats};
 
 /// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
 ///
@@ -477,3 +479,78 @@ where
         }
     }
 }
+
+pub struct WithQuery<'world, 'state, I, Q, F>
+where
+    Q: WorldQuery,
+    F: WorldQuery,
+{
+    iter: I,
+    query: &'world Query<'world, 'state, Q, F>,
+}
+
+impl<'world, 'state, I, Q, F> Iterator for WithQuery<'world, 'state, I, Q, F>
+where
+    I: Iterator<Item = Entity>,
+    Q: WorldQuery,
+    <Q as WorldQueryGats<'world>>::Fetch: ReadOnlyFetch,
+    F: WorldQuery,
+{
+    type Item = ROQueryItem<'world, Q>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // if the entity iterator is done, this iterator is done
+        let mut entity = self.iter.next()?;
+        loop {
+            match self.query.get(entity).ok() {
+                Some(item) => return Some(item),
+                None => {
+                    entity = self.iter.next()?;
+                }
+            }
+        }
+    }
+}
+
+pub trait Joinable {
+    fn join_with<'world, 'state, Q, F>(
+        self,
+        query: &'world Query<'world, 'state, Q, F>,
+    ) -> WithQuery<'world, 'state, Self, Q, F>
+    where
+        Self: Sized,
+        Q: WorldQuery,
+        F: WorldQuery;
+
+    fn join_with_mut<Q, F, CB>(self, query: &mut Query<Q, F>, cb: CB)
+    where
+        Q: WorldQuery,
+        F: WorldQuery,
+        CB: FnMut(QueryItem<Q>);
+}
+
+impl<I: Iterator<Item = Entity>> Joinable for I {
+    fn join_with<'w, 's, Q, F>(
+        self,
+        query: &'w Query<'w, 's, Q, F>,
+    ) -> WithQuery<'w, 's, Self, Q, F>
+    where
+        Q: WorldQuery,
+        F: WorldQuery,
+    {
+        WithQuery { iter: self, query }
+    }
+
+    fn join_with_mut<Q, F, CB>(self, query: &mut Query<Q, F>, mut cb: CB)
+    where
+        Q: WorldQuery,
+        F: WorldQuery,
+        CB: FnMut(QueryItem<Q>),
+    {
+        self.for_each(|entity| {
+            if let Ok(item) = query.get_mut(entity) {
+                cb(item);
+            }
+        });
+    }
+}

crates/bevy_ecs/src/system/mod.rs

@@ -100,7 +100,7 @@ mod tests {
         bundle::Bundles,
         component::{Component, Components},
         entity::{Entities, Entity},
-        query::{Added, Changed, Or, With, Without},
+        query::{Added, Changed, Joinable, Or, With, Without},
         schedule::{Schedule, Stage, SystemStage},
         system::{
             Commands, IntoExclusiveSystem, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query,
@@ -925,4 +925,51 @@ mod tests {
         assert!(entity.contains::<A>());
         assert!(entity.contains::<B>());
     }
+
+    #[test]
+    fn join_with_query() {
+        fn sys(has_a: Query<Entity, With<A>>, has_a_and_b: Query<(&A, &B)>) {
+            assert_eq!(has_a.iter().join_with(&has_a_and_b).count(), 2);
+        }
+
+        let mut system = IntoSystem::into_system(sys);
+        let mut world = World::new();
+        world.spawn().insert_bundle((A, B));
+        world.spawn().insert_bundle((A,));
+        world.spawn().insert_bundle((A, B));
+        world.spawn().insert_bundle((B,));
+
+        system.initialize(&mut world);
+        system.run((), &mut world);
+    }
+
+    #[test]
+    fn join_with_query_mut() {
+        fn sys(has_a: Query<Entity, With<A>>, mut has_w: Query<&mut W<u64>>) {
+            has_a.iter().join_with_mut(&mut has_w, |mut w| {
+                w.0 = 999;
+            });
+        }
+
+        fn check_system(query: Query<(Option<&A>, &W<u64>)>) {
+            for (maybe_a, w) in query.iter() {
+                match maybe_a {
+                    Some(_) => assert_eq!(w.0, 999),
+                    None => assert_eq!(w.0, 0),
+                }
+            }
+        }
+
+        let mut system_sys = IntoSystem::into_system(sys);
+        let mut system_check = IntoSystem::into_system(check_system);
+        let mut world = World::new();
+        world.spawn().insert_bundle((A, W(0u64)));
+        world.spawn().insert_bundle((A, W(0u64)));
+        world.spawn().insert_bundle((W(0u64),));
+
+        system_sys.initialize(&mut world);
+        system_check.initialize(&mut world);
+        system_sys.run((), &mut world);
+        system_check.run((), &mut world);
+    }
 }