diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 25bc9720b2984..312cc58d624a5 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/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, + Q: WorldQuery, + >::Fetch: ReadOnlyFetch, + F: WorldQuery, +{ + type Item = ROQueryItem<'world, Q>; + + fn next(&mut self) -> Option { + // 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(self, query: &mut Query, cb: CB) + where + Q: WorldQuery, + F: WorldQuery, + CB: FnMut(QueryItem); +} + +impl> 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(self, query: &mut Query, mut cb: CB) + where + Q: WorldQuery, + F: WorldQuery, + CB: FnMut(QueryItem), + { + self.for_each(|entity| { + if let Ok(item) = query.get_mut(entity) { + cb(item); + } + }); + } +} diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index 541cd674be935..c69d7c0496f5b 100644 --- a/crates/bevy_ecs/src/system/mod.rs +++ b/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::()); assert!(entity.contains::()); } + + #[test] + fn join_with_query() { + fn sys(has_a: Query>, 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>, mut has_w: Query<&mut W>) { + has_a.iter().join_with_mut(&mut has_w, |mut w| { + w.0 = 999; + }); + } + + fn check_system(query: Query<(Option<&A>, &W)>) { + 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); + } }