forestia/bevy
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Query Joins (#11535)

Browse Source 
# 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:
Mike 2024-03-11 12:07:36 -07:00 committed by GitHub
parent bc1073ee48
commit 9cd3165105
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 282 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
crates/bevy_ecs/src/archetype.rs
View File

@ -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 {

185
crates/bevy_ecs/src/query/state.rs
View File

@ -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);
}
}

98
crates/bevy_ecs/src/system/query.rs
View File

@ -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> {