mirror of
https://github.com/bevyengine/bevy
synced 2024-11-22 12:43:34 +00:00
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 <alice.i.cecile@gmail.com>
This commit is contained in:
parent
bc1073ee48
commit
9cd3165105
3 changed files with 282 additions and 3 deletions
|
@ -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 {
|
||||
|
|
|
@ -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<D: QueryData, F: QueryFilter> QueryState<D, F> {
|
|||
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<D: QueryData, F: QueryFilter> QueryState<D, F> {
|
|||
}
|
||||
}
|
||||
|
||||
/// 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<OtherD: QueryData, NewD: QueryData>(
|
||||
&self,
|
||||
world: &World,
|
||||
other: &QueryState<OtherD>,
|
||||
) -> QueryState<NewD, ()> {
|
||||
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<OtherD, OtherF>,
|
||||
) -> QueryState<NewD, NewF> {
|
||||
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<TableId> =
|
||||
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<ArchetypeId> =
|
||||
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::<NewD>(),
|
||||
filter = std::any::type_name::<NewF>(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<bevy_ecs::query::state::tests::B>) 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::<A>();
|
||||
world.init_component::<B>();
|
||||
let query = QueryState::<&A>::new(&mut world);
|
||||
let _new_query = query.transmute_filtered::<Entity, Changed<B>>(&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<C>>::new(&mut world);
|
||||
let query_2 = QueryState::<&B, Without<C>>::new(&mut world);
|
||||
let mut new_query: QueryState<Entity, ()> = 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<C>>::new(&mut world);
|
||||
let mut new_query: QueryState<Entity, ()> = 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::<C>();
|
||||
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<bevy_ecs::query::state::tests::C>) \
|
||||
attempts to access terms that are not allowed by state \
|
||||
(&bevy_ecs::query::state::tests::A, bevy_ecs::query::filter::Without<bevy_ecs::query::state::tests::C>) \
|
||||
joined with (&bevy_ecs::query::state::tests::B, bevy_ecs::query::filter::Without<bevy_ecs::query::state::tests::C>)."
|
||||
)]
|
||||
fn cannot_join_wrong_filter() {
|
||||
let mut world = World::new();
|
||||
let query_1 = QueryState::<&A, Without<C>>::new(&mut world);
|
||||
let query_2 = QueryState::<&B, Without<C>>::new(&mut world);
|
||||
let _: QueryState<Entity, Changed<C>> = query_1.join_filtered(&world, &query_2);
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1344,7 +1344,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
|
|||
pub fn transmute_lens_filtered<NewD: QueryData, NewF: QueryFilter>(
|
||||
&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::<NewD, NewF>(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<OtherD: QueryData, NewD: QueryData>(
|
||||
&mut self,
|
||||
other: &mut Query<OtherD>,
|
||||
) -> 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<OtherD, OtherF>,
|
||||
) -> 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::<OtherD, OtherF, NewD, NewF>(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> {
|
||||
|
|
Loading…
Reference in a new issue