From 13a9eb67177a280d836f26c210be9786192d0a47 Mon Sep 17 00:00:00 2001 From: Nathan Fenner Date: Fri, 29 Nov 2024 20:26:53 -0500 Subject: [PATCH] all `.include_entity()` helper on query iterators --- crates/bevy_ecs/src/query/fetch.rs | 6 ++ crates/bevy_ecs/src/query/iter.rs | 35 +++++++ crates/bevy_ecs/src/query/world_query.rs | 123 +++++++++++++++++++++++ 3 files changed, 164 insertions(+) diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 77f40555c28ef..3152763d77cf7 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -280,6 +280,10 @@ pub unsafe trait QueryData: WorldQuery { type ReadOnly: ReadOnlyQueryData::State>; } +unsafe impl QueryData for crate::query::IncludeEntity { + type ReadOnly = crate::query::IncludeEntity; +} + /// A [`QueryData`] that is read only. /// /// # Safety @@ -287,6 +291,8 @@ pub unsafe trait QueryData: WorldQuery { /// This must only be implemented for read-only [`QueryData`]'s. pub unsafe trait ReadOnlyQueryData: QueryData {} +unsafe impl ReadOnlyQueryData for crate::query::IncludeEntity {} + /// The item type returned when a [`WorldQuery`] is iterated over pub type QueryItem<'w, Q> = ::Item<'w>; /// 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 bf4c5432546f9..211d6e7c7d942 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1056,6 +1056,41 @@ 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 ExampleComponent; + /// + /// fn example_system(query: Query<&ExampleComponent>) { + /// for (entity, component) in query.iter().include_entity() { + /// println!("entity {:?}", entity); + /// } + /// } + /// ``` + /// + pub fn include_entity(self) -> QueryIter<'w, 's, crate::query::IncludeEntity, F> { + QueryIter { + world: self.world, + tables: self.tables, + archetypes: self.archetypes, + query_state: unsafe { self.query_state.as_transmuted_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> { diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index f2da8b3558f26..c39cd610d09e4 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/crates/bevy_ecs/src/query/world_query.rs @@ -146,6 +146,129 @@ pub unsafe trait WorldQuery { ) -> bool; } +/// A wrapper type which includes the entity in the selected query data, +/// but which has the same underlying state as `D`. +/// +/// Most applications should use `(Entity, D)` instead; but unlike `(Entity, D)`, +/// it is guaranteed that `IncludeEntity` has the same `WorldQuery::State as `D`. +pub struct IncludeEntity(pub D); + +unsafe impl WorldQuery for IncludeEntity { + /// The item is the same as `D`, but also includes the entity. + type Item<'a> = (Entity, D::Item<'a>); + + /// 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 non-`Entity` 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)) + } + + /// 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) + } + + /// Creates a new instance of this 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); + } + + /// Fetch [`Self::Item`](`WorldQuery::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)) + } + + /// 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)),*) => {