diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 2564223972..bd1032f21d 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -333,6 +333,41 @@ pub unsafe trait QueryData: WorldQuery { ) -> Self::Item<'w, 's>; } +/// SAFETY: +/// +/// The `IncludeEntity` 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 QueryData for crate::query::IncludeEntity { + type ReadOnly = crate::query::IncludeEntity; + /// 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 {} +/// SAFETY: +/// +/// The underlying state and access of `IncludeEntity` 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 ReadOnlyQueryData for crate::query::IncludeEntity {} + /// The item type returned when a [`WorldQuery`] is iterated over pub type QueryItem<'w, 's, Q> = ::Item<'w, 's>; /// The read-only variant of the item type returned when a [`QueryData`] is iterated over immutably diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index eb49204434..927da496bf 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -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, F> { + // SAFETY: `IncludeEntity` 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::>(); + + let mut query_a = world.query::<&A>(); + let iter_include_entity = query_a + .iter(&world) + .include_entity() + .map(|(entity, a)| (entity, *a)) + .collect::>(); + + 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::>(); + + 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::>(); + + assert_eq!(iter_plain, iter_include_entity); + } + + { + // With filter + let mut query_entity_a = + world.query_filtered::<(Entity, &A), crate::query::With>(); + let iter_plain = query_entity_a + .iter(&world) + .map(|(entity, a)| (entity, *a)) + .collect::>(); + + let mut query_a = world.query_filtered::<&A, crate::query::With>(); + let iter_include_entity = query_a + .iter(&world) + .include_entity() + .map(|(entity, a)| (entity, *a)) + .collect::>(); + + 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::>(); + + 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::>(); + + assert_eq!(iter_plain, iter_include_entity); + } + } } diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index 1c739927ac..b85a5fb5e3 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/crates/bevy_ecs/src/query/world_query.rs @@ -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` is guaranteed to have identical +/// internal query state to `D` itself. +pub struct IncludeEntity(pub D); + +/// SAFETY: +/// +/// `IncludeEntity` 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 WorldQuery for IncludeEntity { + /// The same underlying state is used for `IncludeEntity` as `D` itself. + type Fetch<'a> = D::Fetch<'a>; + + /// The same underlying state is used for `IncludeEntity` 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) { + 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) { + 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 { + 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)),*) => {