From 9cd3165105398b1858859b4278d30d0079bed912 Mon Sep 17 00:00:00 2001 From: Mike Date: Mon, 11 Mar 2024 12:07:36 -0700 Subject: [PATCH] Query Joins (#11535) # Objective - Add a way to combine 2 queries together in a similar way to `Query::transmute_lens` - Fixes #1658 ## Solution - Use a similar method to query transmute, but take the intersection of matched archetypes between the 2 queries and the union of the accesses to create the new underlying QueryState. --- ## Changelog - Add query joins --------- Co-authored-by: Alice Cecile --- crates/bevy_ecs/src/archetype.rs | 2 +- crates/bevy_ecs/src/query/state.rs | 185 +++++++++++++++++++++++++++- crates/bevy_ecs/src/system/query.rs | 98 ++++++++++++++- 3 files changed, 282 insertions(+), 3 deletions(-) diff --git a/crates/bevy_ecs/src/archetype.rs b/crates/bevy_ecs/src/archetype.rs index baa16e2503baa..6ab5bb2d420de 100644 --- a/crates/bevy_ecs/src/archetype.rs +++ b/crates/bevy_ecs/src/archetype.rs @@ -594,7 +594,7 @@ impl Archetype { /// /// This is used in archetype update methods to limit archetype updates to the /// ones added since the last time the method ran. -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, PartialEq)] pub struct ArchetypeGeneration(ArchetypeId); impl ArchetypeGeneration { diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 8abd77d00d091..1a97e2bd70b4e 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -10,6 +10,7 @@ use crate::{ storage::{SparseSetIndex, TableId}, world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId}, }; +use bevy_utils::tracing::warn; #[cfg(feature = "trace")] use bevy_utils::tracing::Span; use fixedbitset::FixedBitSet; @@ -374,7 +375,11 @@ impl QueryState { NewF::update_component_access(&filter_state, &mut filter_component_access); component_access.extend(&filter_component_access); - assert!(component_access.is_subset(&self.component_access), "Transmuted state for {} attempts to access terms that are not allowed by original state {}.", std::any::type_name::<(NewD, NewF)>(), std::any::type_name::<(D, F)>() ); + assert!( + component_access.is_subset(&self.component_access), + "Transmuted state for {} attempts to access terms that are not allowed by original state {}.", + std::any::type_name::<(NewD, NewF)>(), std::any::type_name::<(D, F)>() + ); QueryState { world_id: self.world_id, @@ -396,6 +401,114 @@ impl QueryState { } } + /// Use this to combine two queries. The data accessed will be the intersection + /// of archetypes included in both queries. This can be useful for accessing a + /// subset of the entities between two queries. + /// + /// You should not call `update_archetypes` on the returned `QueryState` as the result + /// could be unpredictable. You might end up with a mix of archetypes that only matched + /// the original query + archetypes that only match the new `QueryState`. Most of the + /// safe methods on `QueryState` call [`QueryState::update_archetypes`] internally, so + /// this is best used through a `Query`. + /// + /// ## Performance + /// + /// This will have similar performance as constructing a new `QueryState` since much of internal state + /// needs to be reconstructed. But it will be a little faster as it only needs to compare the intersection + /// of matching archetypes rather than iterating over all archetypes. + /// + /// ## Panics + /// + /// Will panic if `NewD` contains accesses not in `Q` or `OtherQ`. + pub fn join( + &self, + world: &World, + other: &QueryState, + ) -> QueryState { + self.join_filtered::<_, (), NewD, ()>(world, other) + } + + /// Use this to combine two queries. The data accessed will be the intersection + /// of archetypes included in both queries. + /// + /// ## Panics + /// + /// Will panic if `NewD` or `NewF` requires accesses not in `Q` or `OtherQ`. + pub fn join_filtered< + OtherD: QueryData, + OtherF: QueryFilter, + NewD: QueryData, + NewF: QueryFilter, + >( + &self, + world: &World, + other: &QueryState, + ) -> QueryState { + if self.world_id != other.world_id { + panic!("Joining queries initialized on different worlds is not allowed."); + } + + let mut component_access = FilteredAccess::default(); + let mut new_fetch_state = NewD::get_state(world) + .expect("Could not create fetch_state, Please initialize all referenced components before transmuting."); + let new_filter_state = NewF::get_state(world) + .expect("Could not create filter_state, Please initialize all referenced components before transmuting."); + + NewD::set_access(&mut new_fetch_state, &self.component_access); + NewD::update_component_access(&new_fetch_state, &mut component_access); + + let mut new_filter_component_access = FilteredAccess::default(); + NewF::update_component_access(&new_filter_state, &mut new_filter_component_access); + + component_access.extend(&new_filter_component_access); + + let mut joined_component_access = self.component_access.clone(); + joined_component_access.extend(&other.component_access); + + assert!( + component_access.is_subset(&joined_component_access), + "Joined state for {} attempts to access terms that are not allowed by state {} joined with {}.", + std::any::type_name::<(NewD, NewF)>(), std::any::type_name::<(D, F)>(), std::any::type_name::<(OtherD, OtherF)>() + ); + + if self.archetype_generation != other.archetype_generation { + warn!("You have tried to join queries with different archetype_generations. This could lead to unpredictable results."); + } + + // take the intersection of the matched ids + let matched_tables: FixedBitSet = self + .matched_tables + .intersection(&other.matched_tables) + .collect(); + let matched_table_ids: Vec = + matched_tables.ones().map(TableId::from_usize).collect(); + let matched_archetypes: FixedBitSet = self + .matched_archetypes + .intersection(&other.matched_archetypes) + .collect(); + let matched_archetype_ids: Vec = + matched_archetypes.ones().map(ArchetypeId::new).collect(); + + QueryState { + world_id: self.world_id, + archetype_generation: self.archetype_generation, + matched_table_ids, + matched_archetype_ids, + fetch_state: new_fetch_state, + filter_state: new_filter_state, + component_access: joined_component_access, + matched_tables, + matched_archetypes, + archetype_component_access: self.archetype_component_access.clone(), + #[cfg(feature = "trace")] + par_iter_span: bevy_utils::tracing::info_span!( + "par_for_each", + query = std::any::type_name::(), + filter = std::any::type_name::(), + ), + } + } + /// Gets the query result for the given [`World`] and [`Entity`]. /// /// This can only be called for read-only queries, see [`Self::get_mut`] for write-queries. @@ -1658,4 +1771,74 @@ mod tests { assert_eq!(entity_a, detection_query.single(&world)); } + + #[test] + #[should_panic( + expected = "Transmuted state for (bevy_ecs::entity::Entity, bevy_ecs::query::filter::Changed) attempts to access terms that are not allowed by original state (&bevy_ecs::query::state::tests::A, ())." + )] + fn cannot_transmute_changed_without_access() { + let mut world = World::new(); + world.init_component::(); + world.init_component::(); + let query = QueryState::<&A>::new(&mut world); + let _new_query = query.transmute_filtered::>(&world); + } + + #[test] + fn join() { + let mut world = World::new(); + world.spawn(A(0)); + world.spawn(B(1)); + let entity_ab = world.spawn((A(2), B(3))).id(); + world.spawn((A(4), B(5), C(6))); + + let query_1 = QueryState::<&A, Without>::new(&mut world); + let query_2 = QueryState::<&B, Without>::new(&mut world); + let mut new_query: QueryState = query_1.join_filtered(&world, &query_2); + + assert_eq!(new_query.single(&world), entity_ab); + } + + #[test] + fn join_with_get() { + let mut world = World::new(); + world.spawn(A(0)); + world.spawn(B(1)); + let entity_ab = world.spawn((A(2), B(3))).id(); + let entity_abc = world.spawn((A(4), B(5), C(6))).id(); + + let query_1 = QueryState::<&A>::new(&mut world); + let query_2 = QueryState::<&B, Without>::new(&mut world); + let mut new_query: QueryState = query_1.join_filtered(&world, &query_2); + + assert!(new_query.get(&world, entity_ab).is_ok()); + // should not be able to get entity with c. + assert!(new_query.get(&world, entity_abc).is_err()); + } + + #[test] + #[should_panic(expected = "Joined state for (&bevy_ecs::query::state::tests::C, ()) \ + attempts to access terms that are not allowed by state \ + (&bevy_ecs::query::state::tests::A, ()) joined with (&bevy_ecs::query::state::tests::B, ()).")] + fn cannot_join_wrong_fetch() { + let mut world = World::new(); + world.init_component::(); + let query_1 = QueryState::<&A>::new(&mut world); + let query_2 = QueryState::<&B>::new(&mut world); + let _query: QueryState<&C> = query_1.join(&world, &query_2); + } + + #[test] + #[should_panic( + expected = "Joined state for (bevy_ecs::entity::Entity, bevy_ecs::query::filter::Changed) \ + attempts to access terms that are not allowed by state \ + (&bevy_ecs::query::state::tests::A, bevy_ecs::query::filter::Without) \ + joined with (&bevy_ecs::query::state::tests::B, bevy_ecs::query::filter::Without)." + )] + fn cannot_join_wrong_filter() { + let mut world = World::new(); + let query_1 = QueryState::<&A, Without>::new(&mut world); + let query_2 = QueryState::<&B, Without>::new(&mut world); + let _: QueryState> = query_1.join_filtered(&world, &query_2); + } } diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index db34e2513e0ea..e34557dbda268 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -1344,7 +1344,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn transmute_lens_filtered( &mut self, ) -> QueryLens<'_, NewD, NewF> { - // SAFETY: There are no other active borrows of data from world + // SAFETY: + // - We have exclusive access to the query + // - `self` has correctly captured it's access + // - Access is checked to be a subset of the query's access when the state is created. let world = unsafe { self.world.world() }; let state = self.state.transmute_filtered::(world); QueryLens { @@ -1359,6 +1362,99 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn as_query_lens(&mut self) -> QueryLens<'_, D> { self.transmute_lens() } + + /// Returns a [`QueryLens`] that can be used to get a query with the combined fetch. + /// + /// For example, this can take a `Query<&A>` and a `Queryy<&B>` and return a `Query<&A, &B>`. + /// The returned query will only return items with both `A` and `B`. Note that since filter + /// are dropped, non-archetypal filters like `Added` and `Changed` will no be respected. + /// To maintain or change filter terms see `Self::join_filtered`. + /// + /// ## Example + /// + /// ```rust + /// # use bevy_ecs::prelude::*; + /// # use bevy_ecs::system::QueryLens; + /// # + /// # #[derive(Component)] + /// # struct Transform; + /// # + /// # #[derive(Component)] + /// # struct Player; + /// # + /// # #[derive(Component)] + /// # struct Enemy; + /// # + /// # let mut world = World::default(); + /// # world.spawn((Transform, Player)); + /// # world.spawn((Transform, Enemy)); + /// + /// fn system( + /// mut transforms: Query<&Transform>, + /// mut players: Query<&Player>, + /// mut enemies: Query<&Enemy> + /// ) { + /// let mut players_transforms: QueryLens<(&Transform, &Player)> = transforms.join(&mut players); + /// for (transform, player) in &players_transforms.query() { + /// // do something with a and b + /// } + /// + /// let mut enemies_transforms: QueryLens<(&Transform, &Enemy)> = transforms.join(&mut enemies); + /// for (transform, enemy) in &enemies_transforms.query() { + /// // do something with a and b + /// } + /// } + /// + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems(system); + /// # schedule.run(&mut world); + /// ``` + /// ## Panics + /// + /// This will panic if `NewD` is not a subset of the union of the original fetch `Q` and `OtherD`. + /// + /// ## Allowed Transmutes + /// + /// Like `transmute_lens` the query terms can be changed with some restrictions. + /// See [`Self::transmute_lens`] for more details. + pub fn join( + &mut self, + other: &mut Query, + ) -> QueryLens<'_, NewD> { + self.join_filtered(other) + } + + /// Equivalent to [`Self::join`] but also includes a [`QueryFilter`] type. + /// + /// Note that the lens with iterate a subset of the original queries tables + /// and archetypes. This means that additional archetypal query terms like + /// `With` and `Without` will not necessarily be respected and non-archetypal + /// terms like `Added` and `Changed` will only be respected if they are in + /// the type signature. + pub fn join_filtered< + OtherD: QueryData, + OtherF: QueryFilter, + NewD: QueryData, + NewF: QueryFilter, + >( + &mut self, + other: &mut Query, + ) -> QueryLens<'_, NewD, NewF> { + // SAFETY: + // - The queries have correctly captured their access. + // - We have exclusive access to both queries. + // - Access for QueryLens is checked when state is created. + let world = unsafe { self.world.world() }; + let state = self + .state + .join_filtered::(world, other.state); + QueryLens { + world: self.world, + state, + last_run: self.last_run, + this_run: self.this_run, + } + } } impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'w Query<'_, 's, D, F> {