From cef438a6e54528234e020291d8b017135ae50ed9 Mon Sep 17 00:00:00 2001 From: Josh Kuhn Date: Wed, 28 Jul 2021 10:59:10 -0700 Subject: [PATCH 1/4] Create with_query extension trait for Iterator --- crates/bevy_ecs/src/query/iter.rs | 64 ++++++++++++++++++++++++++++++- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 25bc9720b2984..6be2e66aed15d 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, WorldQueryGats}; /// An [`Iterator`] over query results of a [`Query`](crate::system::Query). /// @@ -477,3 +479,61 @@ 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, + F: WorldQuery, +{ + type Item = <>::ReadOnlyFetch as Fetch<'world>>::Item; + + 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 WithQueryExt { + fn with_query<'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; +} + +impl WithQueryExt for I +where + I: Iterator, +{ + fn with_query<'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 } + } +} From 016d8d46f5486ec7ddda13b13b7088669fb7b3a3 Mon Sep 17 00:00:00 2001 From: Josh Kuhn Date: Sat, 7 May 2022 20:39:54 -0700 Subject: [PATCH 2/4] Add tests and with_query_mut --- crates/bevy_ecs/src/query/iter.rs | 30 ++++++++++++++----- crates/bevy_ecs/src/system/mod.rs | 49 ++++++++++++++++++++++++++++++- 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 6be2e66aed15d..8bbf0db090207 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -8,7 +8,7 @@ use crate::{ }; use std::{marker::PhantomData, mem::MaybeUninit}; -use super::{QueryFetch, QueryItem, WorldQueryGats}; +use super::{QueryFetch, QueryItem, ROQueryItem, WorldQueryGats}; /// An [`Iterator`] over query results of a [`Query`](crate::system::Query). /// @@ -483,7 +483,7 @@ where pub struct WithQuery<'world, 'state, I, Q, F> where Q: WorldQuery, - F: WorldQuery + F: WorldQuery, { iter: I, query: &'world Query<'world, 'state, Q, F>, @@ -493,9 +493,10 @@ 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 = <>::ReadOnlyFetch as Fetch<'world>>::Item; + type Item = ROQueryItem<'world, Q>; fn next(&mut self) -> Option { // if the entity iterator is done, this iterator is done @@ -520,12 +521,14 @@ pub trait WithQueryExt { Self: Sized, Q: WorldQuery, F: WorldQuery; + fn with_query_mut(self, query: &mut Query, cb: CB) + where + Q: WorldQuery, + F: WorldQuery, + CB: FnMut(QueryItem); } -impl WithQueryExt for I -where - I: Iterator, -{ +impl> WithQueryExt for I { fn with_query<'w, 's, Q, F>( self, query: &'w Query<'w, 's, Q, F>, @@ -536,4 +539,17 @@ where { WithQuery { iter: self, query } } + + fn with_query_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..afe03cbd02c97 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, Or, With, WithQueryExt, 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().with_query(&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().with_query_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); + } } From 9d619fab5db7ab17fbe7cdc85c5ba9ea7d63ffb9 Mon Sep 17 00:00:00 2001 From: Josh Kuhn Date: Sat, 7 May 2022 20:48:08 -0700 Subject: [PATCH 3/4] Rename some things --- crates/bevy_ecs/src/query/iter.rs | 13 +++++++------ crates/bevy_ecs/src/system/mod.rs | 6 +++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 8bbf0db090207..312cc58d624a5 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -512,8 +512,8 @@ where } } -pub trait WithQueryExt { - fn with_query<'world, 'state, Q, F>( +pub trait Joinable { + fn join_with<'world, 'state, Q, F>( self, query: &'world Query<'world, 'state, Q, F>, ) -> WithQuery<'world, 'state, Self, Q, F> @@ -521,15 +521,16 @@ pub trait WithQueryExt { Self: Sized, Q: WorldQuery, F: WorldQuery; - fn with_query_mut(self, query: &mut Query, cb: CB) + + fn join_with_mut(self, query: &mut Query, cb: CB) where Q: WorldQuery, F: WorldQuery, CB: FnMut(QueryItem); } -impl> WithQueryExt for I { - fn with_query<'w, 's, Q, F>( +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> @@ -540,7 +541,7 @@ impl> WithQueryExt for I { WithQuery { iter: self, query } } - fn with_query_mut(self, query: &mut Query, mut cb: CB) + fn join_with_mut(self, query: &mut Query, mut cb: CB) where Q: WorldQuery, F: WorldQuery, diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index afe03cbd02c97..13fdbb1c99e8c 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, WithQueryExt, Without}, + query::{Added, Changed, Or, With, Joinable, Without}, schedule::{Schedule, Stage, SystemStage}, system::{ Commands, IntoExclusiveSystem, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query, @@ -929,7 +929,7 @@ mod tests { #[test] fn join_with_query() { fn sys(has_a: Query>, has_a_and_b: Query<(&A, &B)>) { - assert_eq!(has_a.iter().with_query(&has_a_and_b).count(), 2); + assert_eq!(has_a.iter().join_with(&has_a_and_b).count(), 2); } let mut system = IntoSystem::into_system(sys); @@ -946,7 +946,7 @@ mod tests { #[test] fn join_with_query_mut() { fn sys(has_a: Query>, mut has_w: Query<&mut W>) { - has_a.iter().with_query_mut(&mut has_w, |mut w| { + has_a.iter().join_with_mut(&mut has_w, |mut w| { w.0 = 999; }); } From 241801a1397912c8a526363eab448f15e1253dd6 Mon Sep 17 00:00:00 2001 From: Josh Kuhn Date: Sat, 7 May 2022 20:50:52 -0700 Subject: [PATCH 4/4] Fix formatting --- crates/bevy_ecs/src/system/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index 13fdbb1c99e8c..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, Joinable, Without}, + query::{Added, Changed, Joinable, Or, With, Without}, schedule::{Schedule, Stage, SystemStage}, system::{ Commands, IntoExclusiveSystem, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query,