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

Merge de8212b5ed into 877d278785

Browse Source
This commit is contained in:
Nathan Fenner 2025-07-17 17:08:43 -05:00 committed by GitHub
commit 21a27d9d10
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 286 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
crates/bevy_ecs/src/query/fetch.rs
View File

@ -333,6 +333,41 @@ pub unsafe trait QueryData: WorldQuery {
) -> Self::Item<'w, 's>;
}
/// SAFETY:
///
/// The `IncludeEntity<D>` wrapper type has no additional state over the
/// underlying query type `D`. Its access is identical to the access of
/// `D`, except that the `Entity` being queried is also included in the
/// output.
unsafe impl<D: QueryData> QueryData for crate::query::IncludeEntity<D> {
type ReadOnly = crate::query::IncludeEntity<D::ReadOnly>;
/// The item is the same as `D`, but also includes the entity.
type Item<'a> = (Entity, D::Item<'a>);
/// This function manually implements subtyping for the query items.
/// The non-`Entity` part of the query is shrunk exactly how `D` specifies.
fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> {
(item.0, D::shrink(item.1))
}
/// Fetch [`Self::Item`](`QueryData::Item`) for either the given `entity` in the current [`Table`],
/// or for the given `entity` in the current [`Archetype`]. This must always be called after
/// [`WorldQuery::set_table`] with a `table_row` in the range of the current [`Table`] or after
/// [`WorldQuery::set_archetype`] with a `entity` in the current archetype.
///
/// # Safety
///
/// Must always be called _after_ [`WorldQuery::set_table`] or [`WorldQuery::set_archetype`]. `entity` and
/// `table_row` must be in the range of the current table and archetype.
unsafe fn fetch<'w>(
fetch: &mut Self::Fetch<'w>,
entity: Entity,
table_row: TableRow,
) -> Self::Item<'w> {
(entity, D::fetch(fetch, entity, table_row))
}
}
/// A [`QueryData`] that is read only.
///
/// # Safety
@ -340,6 +375,15 @@ pub unsafe trait QueryData: WorldQuery {
/// This must only be implemented for read-only [`QueryData`]'s.
pub unsafe trait ReadOnlyQueryData: QueryData<ReadOnly = Self> {}
/// SAFETY:
///
/// The underlying state and access of `IncludeEntity<D>` is identical to `D` itself,
/// except that the iterated `Entity` is also included in the output.
///
/// Including the `Entity` in the output does not make any additional kind of mutation
/// possible.
unsafe impl<D: ReadOnlyQueryData> ReadOnlyQueryData for crate::query::IncludeEntity<D> {}
/// The item type returned when a [`WorldQuery`] is iterated over
pub type QueryItem<'w, 's, Q> = <Q as QueryData>::Item<'w, 's>;
/// The read-only variant of the item type returned when a [`QueryData`] is iterated over immutably

136
crates/bevy_ecs/src/query/iter.rs
View File

@ -900,6 +900,56 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
)
}
}
/// Transforms the iterator output to include the `Entity` which is being iterated over.
///
/// # Example
///
/// ```
/// use bevy_ecs::prelude::*;
/// #[derive(Component)]
/// struct Health(f32);
///
/// #[derive(Component)]
/// struct Damage(f32);
///
/// fn apply_damage_system(mut query: Query<(&mut Health, &Damage)>) {
/// // Iterate over the input query:
/// for (mut health, damage) in query.iter_mut() {
/// health.0 -= damage.0;
/// }
///
/// // Iterate again, this time including the entity:
/// for (entity, (health, damage)) in query.iter().include_entity() {
/// if health.0 <= 0.0 {
/// println!("entity {:?} has been reduce to 0 health by {} damage", entity, damage.0);
/// }
/// }
/// }
/// ```
///
pub fn include_entity(self) -> QueryIter<'w, 's, crate::query::IncludeEntity<D>, F> {
// SAFETY: `IncludeEntity<D>` and `D` have identical access and query the same archetypes,
// since the internal state is not affected in any way by the `IncludeEntity` wrapper.
let query_state = unsafe { self.query_state.as_transmuted_state() };
QueryIter {
world: self.world,
tables: self.tables,
archetypes: self.archetypes,
query_state,
cursor: QueryIterationCursor {
is_dense: self.cursor.is_dense,
storage_id_iter: self.cursor.storage_id_iter,
table_entities: self.cursor.table_entities,
archetype_entities: self.cursor.archetype_entities,
fetch: self.cursor.fetch,
filter: self.cursor.filter,
current_len: self.cursor.current_len,
current_row: self.cursor.current_row,
},
}
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> {
@ -3009,4 +3059,90 @@ mod tests {
while query.fetch_next().is_some() {}
}
}
#[test]
fn query_iter_include_entity() {
let mut world = World::new();
world.spawn_batch([A(0.), A(1.), A(2.), A(3.), A(4.)]);
world.spawn_batch([
(A(0.), Sparse(100)),
(A(1.), Sparse(101)),
(A(2.), Sparse(102)),
(A(3.), Sparse(103)),
(A(4.), Sparse(104)),
]);
{
// Read-only
let mut query_entity_a = world.query::<(Entity, &A)>();
let iter_plain = query_entity_a
.iter(&world)
.map(|(entity, a)| (entity, *a))
.collect::<Vec<_>>();
let mut query_a = world.query::<&A>();
let iter_include_entity = query_a
.iter(&world)
.include_entity()
.map(|(entity, a)| (entity, *a))
.collect::<Vec<_>>();
assert_eq!(iter_plain, iter_include_entity);
}
{
// Mut
let mut query_entity_a = world.query::<(Entity, &mut A)>();
let iter_plain = query_entity_a
.iter_mut(&mut world)
.map(|(entity, a)| (entity, *a))
.collect::<Vec<_>>();
let mut query_a = world.query::<&mut A>();
let iter_include_entity = query_a
.iter_mut(&mut world)
.include_entity()
.map(|(entity, a)| (entity, *a))
.collect::<Vec<_>>();
assert_eq!(iter_plain, iter_include_entity);
}
{
// With filter
let mut query_entity_a =
world.query_filtered::<(Entity, &A), crate::query::With<Sparse>>();
let iter_plain = query_entity_a
.iter(&world)
.map(|(entity, a)| (entity, *a))
.collect::<Vec<_>>();
let mut query_a = world.query_filtered::<&A, crate::query::With<Sparse>>();
let iter_include_entity = query_a
.iter(&world)
.include_entity()
.map(|(entity, a)| (entity, *a))
.collect::<Vec<_>>();
assert_eq!(iter_plain, iter_include_entity);
}
{
// Multiple components in a query
let mut query_entity_a = world.query::<(Entity, &A, &Sparse)>();
let iter_plain = query_entity_a
.iter(&world)
.map(|(entity, a, sparse)| (entity, *a, *sparse))
.collect::<Vec<_>>();
let mut query_a = world.query::<(&A, &Sparse)>();
let iter_include_entity = query_a
.iter(&world)
.include_entity()
.map(|(entity, (a, sparse))| (entity, *a, *sparse))
.collect::<Vec<_>>();
assert_eq!(iter_plain, iter_include_entity);
}
}
}

106
crates/bevy_ecs/src/query/world_query.rs
View File

@ -132,6 +132,112 @@ pub unsafe trait WorldQuery {
) -> bool;
}
/// A wrapper type around a data query `D` which will return the queried [`crate::entity::Entity`]
/// alongside the results from `D`.
///
/// This convenience wrapper can be used by calling [`super::iter::QueryIter::include_entity`] on
/// an iterator returned from `[Query::iter()]` or `[Query::iter_mut()]`.
///
/// Unlike `(Entity, D)`, the type `IncludeEntity<D>` is guaranteed to have identical
/// internal query state to `D` itself.
pub struct IncludeEntity<D>(pub D);
/// SAFETY:
///
/// `IncludeEntity<D>` has the exact same internal state and query behavior as `D` itself,
/// and inherits all of its safety properties from this fact.
///
/// The only difference is that the `Entity` (which is tracked by all queries internally) is
/// also included in the output `Item`.
unsafe impl<D: WorldQuery> WorldQuery for IncludeEntity<D> {
/// The same underlying state is used for `IncludeEntity<D>` as `D` itself.
type Fetch<'a> = D::Fetch<'a>;
/// The same underlying state is used for `IncludeEntity<D>` as `D` itself.
type State = D::State;
/// This function manually implements subtyping for the query items.
/// The query is shrunk exactly how `D` specifies.
fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
D::shrink_fetch(fetch)
}
unsafe fn init_fetch<'w>(
world: UnsafeWorldCell<'w>,
state: &Self::State,
last_run: Tick,
this_run: Tick,
) -> Self::Fetch<'w> {
D::init_fetch(world, state, last_run, this_run)
}
const IS_DENSE: bool = D::IS_DENSE;
/// # Safety
///
/// - `archetype` and `tables` must be from the same [`World`] that [`WorldQuery::init_state`] was called on.
/// - `table` must correspond to `archetype`.
/// - `state` must be the [`State`](Self::State) that `fetch` was initialized with.
unsafe fn set_archetype<'w>(
fetch: &mut Self::Fetch<'w>,
state: &Self::State,
archetype: &'w Archetype,
table: &'w Table,
) {
D::set_archetype(fetch, state, archetype, table);
}
/// Adjusts internal state to account for the next [`Table`]. This will always be called on tables
/// that match this [`WorldQuery`].
///
/// # Safety
///
/// - `table` must be from the same [`World`] that [`WorldQuery::init_state`] was called on.
/// - `state` must be the [`State`](Self::State) that `fetch` was initialized with.
unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) {
D::set_table(fetch, state, table);
}
/// Sets available accesses for implementors with dynamic access such as [`FilteredEntityRef`](crate::world::FilteredEntityRef)
/// or [`FilteredEntityMut`](crate::world::FilteredEntityMut).
///
/// Called when constructing a [`QueryLens`](crate::system::QueryLens) or calling [`QueryState::from_builder`](super::QueryState::from_builder)
fn set_access(state: &mut Self::State, access: &FilteredAccess<ComponentId>) {
D::set_access(state, access);
}
/// Adds any component accesses used by this [`WorldQuery`] to `access`.
///
/// Used to check which queries are disjoint and can run in parallel
// This does not have a default body of `{}` because 99% of cases need to add accesses
// and forgetting to do so would be unsound.
fn update_component_access(state: &Self::State, access: &mut FilteredAccess<ComponentId>) {
D::update_component_access(state, access);
}
/// Creates and initializes a [`State`](WorldQuery::State) for this [`WorldQuery`] type.
fn init_state(world: &mut World) -> Self::State {
D::init_state(world)
}
/// Attempts to initialize a [`State`](WorldQuery::State) for this [`WorldQuery`] type using read-only
/// access to [`Components`].
fn get_state(components: &Components) -> Option<Self::State> {
D::get_state(components)
}
/// Returns `true` if this query matches a set of components. Otherwise, returns `false`.
///
/// Used to check which [`Archetype`]s can be skipped by the query
/// (if none of the [`Component`](crate::component::Component)s match)
fn matches_component_set(
state: &Self::State,
set_contains_id: &impl Fn(ComponentId) -> bool,
) -> bool {
D::matches_component_set(state, set_contains_id)
}
}
macro_rules! impl_tuple_world_query {
($(#[$meta:meta])* $(($name: ident, $state: ident)),*) => {