diff --git a/crates/bevy_asset/src/asset_changed.rs b/crates/bevy_asset/src/asset_changed.rs index 555861dd9809f..9721a3d1d6b2f 100644 --- a/crates/bevy_asset/src/asset_changed.rs +++ b/crates/bevy_asset/src/asset_changed.rs @@ -201,6 +201,7 @@ unsafe impl WorldQuery for AssetChanged { } const IS_DENSE: bool = <&A>::IS_DENSE; + const IS_ARCHETYPAL: bool = false; unsafe fn set_archetype<'w, 's>( fetch: &mut Self::Fetch<'w>, @@ -261,30 +262,26 @@ unsafe impl WorldQuery for AssetChanged { ) -> bool { set_contains_id(state.asset_id) } -} - -#[expect(unsafe_code, reason = "QueryFilter is an unsafe trait.")] -/// SAFETY: read-only access -unsafe impl QueryFilter for AssetChanged { - const IS_ARCHETYPAL: bool = false; #[inline] - unsafe fn filter_fetch( + unsafe fn matches( state: &Self::State, fetch: &mut Self::Fetch<'_>, entity: Entity, table_row: TableRow, ) -> bool { fetch.inner.as_mut().is_some_and(|inner| { - // SAFETY: We delegate to the inner `fetch` for `A` + // SAFETY: We delegate to the inner `fetch` for `A`. unsafe { - let handle = <&A>::fetch(&state.asset_id, inner, entity, table_row); + let handle = <&A>::try_fetch(&state.asset_id, inner, entity, table_row); handle.is_some_and(|handle| fetch.check.has_changed(handle)) } }) } } +impl QueryFilter for AssetChanged {} + #[cfg(test)] #[expect(clippy::print_stdout, reason = "Allowed in tests.")] mod tests { diff --git a/crates/bevy_ecs/macros/src/query_data.rs b/crates/bevy_ecs/macros/src/query_data.rs index a91dd193e51c1..1b66cdde2c7f2 100644 --- a/crates/bevy_ecs/macros/src/query_data.rs +++ b/crates/bevy_ecs/macros/src/query_data.rs @@ -233,7 +233,6 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { unsafe impl #user_impl_generics #path::query::QueryData for #read_only_struct_name #user_ty_generics #user_where_clauses { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true #(&& <#read_only_field_types as #path::query::QueryData>::IS_ARCHETYPAL)*; type ReadOnly = #read_only_struct_name #user_ty_generics; type Item<'__w, '__s> = #read_only_item_struct_name #user_ty_generics_with_world_and_state; @@ -262,10 +261,10 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { _fetch: &mut ::Fetch<'__w>, _entity: #path::entity::Entity, _table_row: #path::storage::TableRow, - ) -> Option> { - Some(Self::Item { - #(#field_members: <#read_only_field_types>::fetch(&_state.#field_aliases, &mut _fetch.#field_aliases, _entity, _table_row)?,)* - }) + ) -> Self::Item<'__w, '__s> { + Self::Item { + #(#field_members: <#read_only_field_types>::fetch(&_state.#field_aliases, &mut _fetch.#field_aliases, _entity, _table_row),)* + } } fn iter_access( @@ -304,7 +303,6 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { unsafe impl #user_impl_generics #path::query::QueryData for #struct_name #user_ty_generics #user_where_clauses { const IS_READ_ONLY: bool = #is_read_only; - const IS_ARCHETYPAL: bool = true #(&& <#field_types as #path::query::QueryData>::IS_ARCHETYPAL)*; type ReadOnly = #read_only_struct_name #user_ty_generics; type Item<'__w, '__s> = #item_struct_name #user_ty_generics_with_world_and_state; @@ -333,10 +331,10 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { _fetch: &mut ::Fetch<'__w>, _entity: #path::entity::Entity, _table_row: #path::storage::TableRow, - ) -> Option> { - Some(Self::Item { - #(#field_members: <#field_types>::fetch(&_state.#field_aliases, &mut _fetch.#field_aliases, _entity, _table_row)?,)* - }) + ) -> Self::Item<'__w, '__s> { + Self::Item { + #(#field_members: <#field_types>::fetch(&_state.#field_aliases, &mut _fetch.#field_aliases, _entity, _table_row),)* + } } fn iter_access( diff --git a/crates/bevy_ecs/macros/src/query_filter.rs b/crates/bevy_ecs/macros/src/query_filter.rs index ed4f55c6af02f..a7f44ebc53592 100644 --- a/crates/bevy_ecs/macros/src/query_filter.rs +++ b/crates/bevy_ecs/macros/src/query_filter.rs @@ -71,22 +71,7 @@ pub fn derive_query_filter_impl(input: TokenStream) -> TokenStream { ); let filter_impl = quote! { - // SAFETY: This only performs access that subqueries perform, and they impl `QueryFilter` and so perform no mutable access. - unsafe impl #user_impl_generics #path::query::QueryFilter - for #struct_name #user_ty_generics #user_where_clauses { - const IS_ARCHETYPAL: bool = true #(&& <#field_types as #path::query::QueryFilter>::IS_ARCHETYPAL)*; - - #[allow(unused_variables)] - #[inline(always)] - unsafe fn filter_fetch<'__w>( - _state: &Self::State, - _fetch: &mut ::Fetch<'__w>, - _entity: #path::entity::Entity, - _table_row: #path::storage::TableRow, - ) -> bool { - true #(&& <#field_types>::filter_fetch(&_state.#field_aliases, &mut _fetch.#field_aliases, _entity, _table_row))* - } - } + impl #user_impl_generics #path::query::QueryFilter for #struct_name #user_ty_generics #user_where_clauses {} }; let filter_asserts = quote! { diff --git a/crates/bevy_ecs/macros/src/world_query.rs b/crates/bevy_ecs/macros/src/world_query.rs index 39feecf3798b8..874238c84d5eb 100644 --- a/crates/bevy_ecs/macros/src/world_query.rs +++ b/crates/bevy_ecs/macros/src/world_query.rs @@ -130,6 +130,7 @@ pub(crate) fn world_query_impl( } const IS_DENSE: bool = true #(&& <#field_types>::IS_DENSE)*; + const IS_ARCHETYPAL: bool = true #(&& <#field_types>::IS_ARCHETYPAL)*; /// SAFETY: we call `set_archetype` for each member that implements `Fetch` #[inline] @@ -171,6 +172,57 @@ pub(crate) fn world_query_impl( fn matches_component_set(state: &Self::State, _set_contains_id: &impl Fn(#path::component::ComponentId) -> bool) -> bool { true #(&& <#field_types>::matches_component_set(&state.#field_aliases, _set_contains_id))* } + + #[inline] + unsafe fn find_table_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + table_entities: &[#path::entity::Entity], + mut rows: core::ops::Range, + ) -> core::ops::Range { + if Self::IS_ARCHETYPAL { + rows + } else { + // SAFETY: `rows` is only ever narrowed as we iterate subqueries, so it's + // always valid to pass to the next term. Other invariants are upheld by + // the caller. + #(rows = unsafe { <#field_types>::find_table_chunk(&state.#field_aliases, &mut fetch.#field_aliases, table_entities, rows) };)* + rows + } + } + + #[inline] + unsafe fn find_archetype_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + archetype_entities: &[#path::archetype::ArchetypeEntity], + mut indices: core::ops::Range, + ) -> core::ops::Range { + if Self::IS_ARCHETYPAL { + indices + } else { + // SAFETY: `indices` is only ever narrowed as we iterate subqueries, so it's + // always valid to pass to the next term. Other invariants are upheld by + // the caller. + #(indices = unsafe { <#field_types>::find_archetype_chunk(&state.#field_aliases, &mut fetch.#field_aliases, archetype_entities, indices) };)* + indices + } + } + + #[inline] + unsafe fn matches( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + entity: #path::entity::Entity, + table_row: #path::storage::TableRow, + ) -> bool { + if Self::IS_ARCHETYPAL { + true + } else { + // SAFETY: invariants are upheld by the caller. + true #(&& unsafe { <#field_types>::matches(&state.#field_aliases, &mut fetch.#field_aliases, entity, table_row) })* + } + } } } } diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index d0b191513e8ed..83e391c2e0d64 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -1,12 +1,12 @@ use crate::{ - archetype::{Archetype, Archetypes}, + archetype::{Archetype, ArchetypeEntity, Archetypes}, bundle::Bundle, change_detection::{ComponentTicksMut, ComponentTicksRef, MaybeLocation, Tick}, component::{Component, ComponentId, Components, Mutable, StorageType}, entity::{Entities, Entity, EntityLocation}, query::{ access_iter::{EcsAccessLevel, EcsAccessType}, - Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery, + Access, DebugCheckedUnwrap, FilteredAccess, RangeExt, WorldQuery, }, storage::{ComponentSparseSet, Table, TableRow}, world::{ @@ -16,7 +16,7 @@ use crate::{ }; use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref}; use bevy_utils::prelude::DebugName; -use core::{cell::UnsafeCell, iter, marker::PhantomData, panic::Location}; +use core::{cell::UnsafeCell, iter, marker::PhantomData, ops::Range, panic::Location}; use variadics_please::all_tuples; /// Types that can be fetched from a [`World`] using a [`Query`]. @@ -287,15 +287,6 @@ pub unsafe trait QueryData: WorldQuery { /// True if this query is read-only and may not perform mutable access. const IS_READ_ONLY: bool; - /// Returns true if (and only if) this query data relies strictly on archetypes to limit which - /// entities are accessed by the Query. - /// - /// This enables optimizations for [`QueryIter`](`crate::query::QueryIter`) that rely on knowing exactly how - /// many elements are being iterated (such as `Iterator::collect()`). - /// - /// If this is `true`, then [`QueryData::fetch`] must always return `Some`. - const IS_ARCHETYPAL: bool; - /// The read-only variant of this [`QueryData`], which satisfies the [`ReadOnlyQueryData`] trait. type ReadOnly: ReadOnlyQueryData::State>; @@ -331,20 +322,50 @@ pub unsafe trait QueryData: WorldQuery { /// [`WorldQuery::set_archetype`] with an `entity` in the current archetype. /// Accesses components registered in [`WorldQuery::update_component_access`]. /// - /// This method returns `None` if the entity does not match the query. - /// If `Self` implements [`ArchetypeQueryData`], this must always return `Some`. - /// /// # 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. + /// - Must always be called _after_ [`WorldQuery::matches`] if and only if it returns `true`, + /// or after one of [`WorldQuery::find_table_chunk`] or [`WorldQuery::find_archetype_chunk`] + /// and the provided table row/index is within the range returned. /// - There must not be simultaneous conflicting component access registered in `update_component_access`. unsafe fn fetch<'w, 's>( state: &'s Self::State, fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option>; + ) -> Self::Item<'w, 's>; + + /// Attempt to 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 an `entity` in the current archetype. + /// Accesses components registered in [`WorldQuery::update_component_access`]. + /// + /// Unlike [`Self::fetch`](`QueryData::fetch`), this method is fallible, so it is not a requirement + /// to check that the passed row is valid for this query beforehand. This method returns `None` if + /// the entity does not match the query. This method must return `Some()` if and only if + /// [`Self::matches`](`WorldQuery::matches`) would return `true` for the given row, and in that + /// case must return the same item that [`Self::fetch`](`QueryData::fetch`) would return. + /// + /// # 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. + /// - There must not be simultaneous conflicting component access registered in `update_component_access`. + #[inline(always)] + unsafe fn try_fetch<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entity: Entity, + table_row: TableRow, + ) -> Option> { + Self::matches(state, fetch, entity, table_row).then(|| { + // SAFETY: this is only called if the given row is valid for this query + unsafe { Self::fetch(state, fetch, entity, table_row) } + }) + } /// Returns an iterator over the access needed by [`QueryData::fetch`]. Access conflicts are usually /// checked in [`WorldQuery::update_component_access`], but in certain cases this method can be useful to implement @@ -380,7 +401,7 @@ pub trait ReleaseStateQueryData: QueryData { /// This is needed to implement [`ExactSizeIterator`] for /// [`QueryIter`](crate::query::QueryIter) that contains archetype-level filters. /// -/// The trait must only be implemented for query data where its corresponding [`QueryData::IS_ARCHETYPAL`] is [`prim@true`]. +/// The trait must only be implemented for query data where its corresponding [`WorldQuery::IS_ARCHETYPAL`] is [`prim@true`]. pub trait ArchetypeQueryData: QueryData {} /// SAFETY: @@ -401,6 +422,7 @@ unsafe impl WorldQuery for Entity { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -433,12 +455,21 @@ unsafe impl WorldQuery for Entity { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for Entity { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = Entity; @@ -455,8 +486,8 @@ unsafe impl QueryData for Entity { _fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { - Some(entity) + ) -> Self::Item<'w, 's> { + entity } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -498,6 +529,7 @@ unsafe impl WorldQuery for EntityLocation { // This is set to true to avoid forcing archetypal iteration in compound queries, is likely to be slower // in most practical use case. const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -530,12 +562,21 @@ unsafe impl WorldQuery for EntityLocation { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for EntityLocation { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = EntityLocation; @@ -551,9 +592,9 @@ unsafe impl QueryData for EntityLocation { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: `fetch` must be called with an entity that exists in the world - Some(unsafe { fetch.get_spawned(entity).debug_checked_unwrap() }) + unsafe { fetch.get_spawned(entity).debug_checked_unwrap() } } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -666,6 +707,7 @@ unsafe impl WorldQuery for SpawnDetails { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -698,6 +740,16 @@ unsafe impl WorldQuery for SpawnDetails { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } // SAFETY: @@ -705,7 +757,6 @@ unsafe impl WorldQuery for SpawnDetails { // Is its own ReadOnlyQueryData. unsafe impl QueryData for SpawnDetails { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = Self; @@ -721,19 +772,19 @@ unsafe impl QueryData for SpawnDetails { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: only living entities are queried let (spawned_by, spawn_tick) = unsafe { fetch .entities .entity_get_spawned_or_despawned_unchecked(entity) }; - Some(Self { + Self { spawned_by, spawn_tick, last_run: fetch.last_run, this_run: fetch.this_run, - }) + } } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -788,6 +839,7 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -826,12 +878,21 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl<'a> QueryData for EntityRef<'a> { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = EntityRef<'w>; @@ -847,7 +908,7 @@ unsafe impl<'a> QueryData for EntityRef<'a> { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: `fetch` must be called with an entity that exists in the world let cell = unsafe { fetch @@ -856,7 +917,7 @@ unsafe impl<'a> QueryData for EntityRef<'a> { .debug_checked_unwrap() }; // SAFETY: Read-only access to every component has been registered. - Some(unsafe { EntityRef::new(cell) }) + unsafe { EntityRef::new(cell) } } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -898,6 +959,7 @@ unsafe impl<'a> WorldQuery for EntityMut<'a> { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -936,12 +998,21 @@ unsafe impl<'a> WorldQuery for EntityMut<'a> { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: access of `EntityRef` is a subset of `EntityMut` unsafe impl<'a> QueryData for EntityMut<'a> { const IS_READ_ONLY: bool = false; - const IS_ARCHETYPAL: bool = true; type ReadOnly = EntityRef<'a>; type Item<'w, 's> = EntityMut<'w>; @@ -957,7 +1028,7 @@ unsafe impl<'a> QueryData for EntityMut<'a> { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: `fetch` must be called with an entity that exists in the world let cell = unsafe { fetch @@ -966,7 +1037,7 @@ unsafe impl<'a> QueryData for EntityMut<'a> { .debug_checked_unwrap() }; // SAFETY: mutable access to every component has been registered. - Some(unsafe { EntityMut::new(cell) }) + unsafe { EntityMut::new(cell) } } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -992,6 +1063,7 @@ unsafe impl WorldQuery for FilteredEntityRef<'_, '_> { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; unsafe fn init_fetch<'w, 's>( world: UnsafeWorldCell<'w>, @@ -1045,12 +1117,21 @@ unsafe impl WorldQuery for FilteredEntityRef<'_, '_> { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl<'a, 'b> QueryData for FilteredEntityRef<'a, 'b> { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = FilteredEntityRef<'w, 's>; @@ -1085,7 +1166,7 @@ unsafe impl<'a, 'b> QueryData for FilteredEntityRef<'a, 'b> { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: `fetch` must be called with an entity that exists in the world let cell = unsafe { fetch @@ -1094,7 +1175,7 @@ unsafe impl<'a, 'b> QueryData for FilteredEntityRef<'a, 'b> { .debug_checked_unwrap() }; // SAFETY: mutable access to every component has been registered. - Some(unsafe { FilteredEntityRef::new(cell, access) }) + unsafe { FilteredEntityRef::new(cell, access) } } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -1117,6 +1198,7 @@ unsafe impl WorldQuery for FilteredEntityMut<'_, '_> { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; unsafe fn init_fetch<'w, 's>( world: UnsafeWorldCell<'w>, @@ -1170,12 +1252,21 @@ unsafe impl WorldQuery for FilteredEntityMut<'_, '_> { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: access of `FilteredEntityRef` is a subset of `FilteredEntityMut` unsafe impl<'a, 'b> QueryData for FilteredEntityMut<'a, 'b> { const IS_READ_ONLY: bool = false; - const IS_ARCHETYPAL: bool = true; type ReadOnly = FilteredEntityRef<'a, 'b>; type Item<'w, 's> = FilteredEntityMut<'w, 's>; @@ -1208,7 +1299,7 @@ unsafe impl<'a, 'b> QueryData for FilteredEntityMut<'a, 'b> { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: `fetch` must be called with an entity that exists in the world let cell = unsafe { fetch @@ -1217,7 +1308,7 @@ unsafe impl<'a, 'b> QueryData for FilteredEntityMut<'a, 'b> { .debug_checked_unwrap() }; // SAFETY: mutable access to every component has been registered. - Some(unsafe { FilteredEntityMut::new(cell, access) }) + unsafe { FilteredEntityMut::new(cell, access) } } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -1255,6 +1346,7 @@ where } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; unsafe fn set_archetype<'w, 's>( _: &mut Self::Fetch<'w>, @@ -1303,6 +1395,16 @@ where fn matches_component_set(_: &Self::State, _: &impl Fn(ComponentId) -> bool) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly`. @@ -1311,7 +1413,6 @@ where B: Bundle, { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = EntityRefExcept<'w, 's, B>; @@ -1326,12 +1427,12 @@ where fetch: &mut Self::Fetch<'w>, entity: Entity, _: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { let cell = fetch .world .get_entity_with_ticks(entity, fetch.last_run, fetch.this_run) .unwrap(); - Some(EntityRefExcept::new(cell, access)) + EntityRefExcept::new(cell, access) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -1373,6 +1474,7 @@ where } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; unsafe fn set_archetype<'w, 's>( _: &mut Self::Fetch<'w>, @@ -1421,6 +1523,16 @@ where fn matches_component_set(_: &Self::State, _: &impl Fn(ComponentId) -> bool) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: All accesses that `EntityRefExcept` provides are also accesses that @@ -1430,7 +1542,6 @@ where B: Bundle, { const IS_READ_ONLY: bool = false; - const IS_ARCHETYPAL: bool = true; type ReadOnly = EntityRefExcept<'a, 'b, B>; type Item<'w, 's> = EntityMutExcept<'w, 's, B>; @@ -1445,12 +1556,12 @@ where fetch: &mut Self::Fetch<'w>, entity: Entity, _: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { let cell = fetch .world .get_entity_with_ticks(entity, fetch.last_run, fetch.this_run) .unwrap(); - Some(EntityMutExcept::new(cell, access)) + EntityMutExcept::new(cell, access) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -1483,6 +1594,7 @@ unsafe impl WorldQuery for &Archetype { // This could probably be a non-dense query and just set a Option<&Archetype> fetch value in // set_archetypes, but forcing archetypal iteration is likely to be slower in any compound query. const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -1515,12 +1627,21 @@ unsafe impl WorldQuery for &Archetype { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for &Archetype { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = &'w Archetype; @@ -1536,12 +1657,12 @@ unsafe impl QueryData for &Archetype { fetch: &mut Self::Fetch<'w>, entity: Entity, _table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { let (entities, archetypes) = *fetch; // SAFETY: `fetch` must be called with an entity that exists in the world let location = unsafe { entities.get_spawned(entity).debug_checked_unwrap() }; // SAFETY: The assigned archetype for a living entity must always be valid. - Some(unsafe { archetypes.get(location.archetype_id).debug_checked_unwrap() }) + unsafe { archetypes.get(location.archetype_id).debug_checked_unwrap() } } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -1620,6 +1741,8 @@ unsafe impl WorldQuery for &T { } }; + const IS_ARCHETYPAL: bool = true; + #[inline] unsafe fn set_archetype<'w>( fetch: &mut ReadFetch<'w, T>, @@ -1674,12 +1797,21 @@ unsafe impl WorldQuery for &T { ) -> bool { set_contains_id(state) } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for &T { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = &'w T; @@ -1695,8 +1827,8 @@ unsafe impl QueryData for &T { fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { - Some(fetch.components.extract( + ) -> Self::Item<'w, 's> { + fetch.components.extract( |table| { // SAFETY: set_table was previously called let table = unsafe { table.debug_checked_unwrap() }; @@ -1714,7 +1846,7 @@ unsafe impl QueryData for &T { }; item.deref() }, - )) + ) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -1802,6 +1934,7 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w>( @@ -1864,12 +1997,21 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> { ) -> bool { set_contains_id(state) } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl<'__w, T: Component> QueryData for Ref<'__w, T> { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = Ref<'w, T>; @@ -1885,8 +2027,8 @@ unsafe impl<'__w, T: Component> QueryData for Ref<'__w, T> { fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { - Some(fetch.components.extract( + ) -> Self::Item<'w, 's> { + fetch.components.extract( |table| { // SAFETY: set_table was previously called let (table_components, added_ticks, changed_ticks, callers) = @@ -1931,7 +2073,7 @@ unsafe impl<'__w, T: Component> QueryData for Ref<'__w, T> { ), } }, - )) + ) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -2019,6 +2161,7 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w>( @@ -2081,12 +2224,21 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T { ) -> bool { set_contains_id(state) } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: access of `&T` is a subset of `&mut T` unsafe impl<'__w, T: Component> QueryData for &'__w mut T { const IS_READ_ONLY: bool = false; - const IS_ARCHETYPAL: bool = true; type ReadOnly = &'__w T; type Item<'w, 's> = Mut<'w, T>; @@ -2102,8 +2254,8 @@ unsafe impl<'__w, T: Component> QueryData for &'__w mut T fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { - Some(fetch.components.extract( + ) -> Self::Item<'w, 's> { + fetch.components.extract( |table| { // SAFETY: set_table was previously called let (table_components, added_ticks, changed_ticks, callers) = @@ -2148,7 +2300,7 @@ unsafe impl<'__w, T: Component> QueryData for &'__w mut T ), } }, - )) + ) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -2194,6 +2346,7 @@ unsafe impl<'__w, T: Component> WorldQuery for Mut<'__w, T> { // Forwarded to `&mut T` const IS_DENSE: bool = <&mut T as WorldQuery>::IS_DENSE; + const IS_ARCHETYPAL: bool = true; #[inline] // Forwarded to `&mut T` @@ -2241,12 +2394,21 @@ unsafe impl<'__w, T: Component> WorldQuery for Mut<'__w, T> { ) -> bool { <&mut T as WorldQuery>::matches_component_set(state, set_contains_id) } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } // SAFETY: access of `Ref` is a subset of `Mut` unsafe impl<'__w, T: Component> QueryData for Mut<'__w, T> { const IS_READ_ONLY: bool = false; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Ref<'__w, T>; type Item<'w, 's> = Mut<'w, T>; @@ -2266,7 +2428,7 @@ unsafe impl<'__w, T: Component> QueryData for Mut<'__w, T> fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { <&mut T as QueryData>::fetch(state, fetch, entity, table_row) } @@ -2328,6 +2490,9 @@ unsafe impl WorldQuery for Option { } const IS_DENSE: bool = T::IS_DENSE; + // `Option` matches all entities, even if `T` does not, + // so it's always an `ArchetypeQueryData`, even for non-archetypal `T`. + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -2389,14 +2554,21 @@ unsafe impl WorldQuery for Option { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } // SAFETY: defers to soundness of `T: WorldQuery` impl unsafe impl QueryData for Option { const IS_READ_ONLY: bool = T::IS_READ_ONLY; - // `Option` matches all entities, even if `T` does not, - // so it's always an `ArchetypeQueryData`, even for non-archetypal `T`. - const IS_ARCHETYPAL: bool = true; type ReadOnly = Option; type Item<'w, 's> = Option>; @@ -2412,14 +2584,11 @@ unsafe impl QueryData for Option { fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { - Some( - fetch - .matches - // SAFETY: The invariants are upheld by the caller. - .then(|| unsafe { T::fetch(state, &mut fetch.fetch, entity, table_row) }) - .flatten(), - ) + ) -> Self::Item<'w, 's> { + fetch + .matches + // SAFETY: The invariants are upheld by the caller. + .then(|| unsafe { T::fetch(state, &mut fetch.fetch, entity, table_row) }) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -2538,6 +2707,7 @@ unsafe impl WorldQuery for Has { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -2577,12 +2747,21 @@ unsafe impl WorldQuery for Has { // `Has` always matches true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for Has { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = bool; @@ -2598,8 +2777,8 @@ unsafe impl QueryData for Has { fetch: &mut Self::Fetch<'w>, _entity: Entity, _table_row: TableRow, - ) -> Option> { - Some(*fetch) + ) -> Self::Item<'w, 's> { + *fetch } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -2625,6 +2804,13 @@ impl ArchetypeQueryData for Has {} /// Entities are guaranteed to have at least one of the components in `T`. pub struct AnyOf(PhantomData); +#[doc(hidden)] +#[derive(Clone)] +pub struct AnyOfFetch { + fetch: T, + matches: bool, +} + macro_rules! impl_tuple_query_data { ($(#[$meta:meta])* $(($name: ident, $item: ident, $state: ident)),*) => { #[expect( @@ -2647,7 +2833,6 @@ macro_rules! impl_tuple_query_data { // SAFETY: defers to soundness `$name: WorldQuery` impl unsafe impl<$($name: QueryData),*> QueryData for ($($name,)*) { const IS_READ_ONLY: bool = true $(&& $name::IS_READ_ONLY)*; - const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*; type ReadOnly = ($($name::ReadOnly,)*); type Item<'w, 's> = ($($name::Item<'w, 's>,)*); @@ -2674,11 +2859,24 @@ macro_rules! impl_tuple_query_data { fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow + ) -> Self::Item<'w, 's> { + let ($($state,)*) = state; + let ($($name,)*) = fetch; + // SAFETY: The invariants are upheld by the caller. + ($(unsafe { $name::fetch($state, &mut $name.fetch, entity, table_row) },)*) + } + + #[inline(always)] + unsafe fn try_fetch<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entity: Entity, + table_row: TableRow ) -> Option> { let ($($state,)*) = state; let ($($name,)*) = fetch; // SAFETY: The invariants are upheld by the caller. - Some(($(unsafe { $name::fetch($state, $name, entity, table_row) }?,)*)) + Some(($(unsafe { $name::try_fetch($state, &mut $name.fetch, entity, table_row) }?,)*)) } fn iter_access(state: &Self::State) -> impl Iterator> { @@ -2730,30 +2928,47 @@ macro_rules! impl_anytuple_fetch { clippy::unused_unit, reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." )] + #[allow( + unused_mut, + reason = "Zero-length tuples won't mutate any of the parameters." + )] /// SAFETY: /// `fetch` accesses are a subset of the subqueries' accesses /// This is sound because `update_component_access` adds accesses according to the implementations of all the subqueries. /// `update_component_access` replaces the filters with a disjunction where every element is a conjunction of the previous filters and the filters of one of the subqueries. /// This is sound because `matches_component_set` returns a disjunction of the results of the subqueries' implementations. unsafe impl<$($name: WorldQuery),*> WorldQuery for AnyOf<($($name,)*)> { - type Fetch<'w> = ($(($name::Fetch<'w>, bool),)*); + type Fetch<'w> = AnyOfFetch<($(AnyOfFetch<$name::Fetch<'w>>,)*)>; type State = ($($name::State,)*); fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { - let ($($name,)*) = fetch; - ($( - ($name::shrink_fetch($name.0), $name.1), - )*) + let ($($name,)*) = fetch.fetch; + AnyOfFetch { + fetch: + ($(AnyOfFetch { + fetch: $name::shrink_fetch($name.fetch), + matches: $name.matches, + },)*), + matches: fetch.matches, + } } #[inline] unsafe fn init_fetch<'w, 's>(_world: UnsafeWorldCell<'w>, state: &'s Self::State, _last_run: Tick, _this_run: Tick) -> Self::Fetch<'w> { let ($($name,)*) = state; - // SAFETY: The invariants are upheld by the caller. - ($(( unsafe { $name::init_fetch(_world, $name, _last_run, _this_run) }, false),)*) + AnyOfFetch { + fetch: + ($(AnyOfFetch { + // SAFETY: The invariants are upheld by the caller. + fetch: unsafe { $name::init_fetch(_world, $name, _last_run, _this_run) }, + matches: false, + },)*), + matches: false, + } } const IS_DENSE: bool = true $(&& $name::IS_DENSE)*; + const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*; #[inline] unsafe fn set_archetype<'w, 's>( @@ -2762,26 +2977,30 @@ macro_rules! impl_anytuple_fetch { _archetype: &'w Archetype, _table: &'w Table ) { - let ($($name,)*) = _fetch; + let ($($name,)*) = &mut _fetch.fetch; let ($($state,)*) = _state; + _fetch.matches = false; $( - $name.1 = $name::matches_component_set($state, &|id| _archetype.contains(id)); - if $name.1 { + $name.matches = $name::matches_component_set($state, &|id| _archetype.contains(id)); + _fetch.matches = _fetch.matches || $name.matches; + if $name.matches { // SAFETY: The invariants are upheld by the caller. - unsafe { $name::set_archetype(&mut $name.0, $state, _archetype, _table); } + unsafe { $name::set_archetype(&mut $name.fetch, $state, _archetype, _table); } } )* } #[inline] unsafe fn set_table<'w, 's>(_fetch: &mut Self::Fetch<'w>, _state: &'s Self::State, _table: &'w Table) { - let ($($name,)*) = _fetch; + let ($($name,)*) = &mut _fetch.fetch; let ($($state,)*) = _state; + _fetch.matches = false; $( - $name.1 = $name::matches_component_set($state, &|id| _table.has_column(id)); - if $name.1 { + $name.matches = $name::matches_component_set($state, &|id| _table.has_column(id)); + _fetch.matches = _fetch.matches || $name.matches; + if $name.matches { // SAFETY: The invariants are required to be upheld by the caller. - unsafe { $name::set_table(&mut $name.0, $state, _table); } + unsafe { $name::set_table(&mut $name.fetch, $state, _table); } } )* } @@ -2823,6 +3042,81 @@ macro_rules! impl_anytuple_fetch { let ($($name,)*) = _state; false $(|| $name::matches_component_set($name, _set_contains_id))* } + + #[inline] + unsafe fn find_table_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + table_entities: &[Entity], + rows: Range, + ) -> Range { + // If this is an archetypal query, it must match all entities. + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + if Self::IS_ARCHETYPAL || !fetch.matches { + rows + } else { + let ($($name,)*) = &mut fetch.fetch; + let ($($state,)*) = state; + let mut next_chunk = rows.end..rows.end; + $( + if $name.matches { + // SAFETY: invariants are upheld by the caller. + next_chunk = next_chunk.union_or_first(unsafe { $name::find_table_chunk($state, &mut $name.fetch, table_entities, rows.clone()) }); + } + )* + next_chunk + } + } + + #[inline] + unsafe fn find_archetype_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + archetype_entities: &[ArchetypeEntity], + mut indices: Range, + ) -> Range { + // If this is an archetypal query, it must match all entities. + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + if Self::IS_ARCHETYPAL || !fetch.matches { + indices + } else { + let ($($name,)*) = &mut fetch.fetch; + let ($($state,)*) = state; + let mut next_chunk = indices.end..indices.end; + $( + if $name.matches { + // SAFETY: invariants are upheld by the caller. + next_chunk = next_chunk.union_or_first(unsafe { $name::find_archetype_chunk($state, &mut $name.fetch, archetype_entities, indices.clone()) }); + } + )* + next_chunk + } + } + + #[inline] + unsafe fn matches( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + entity: Entity, + table_row: TableRow, + ) -> bool { + // If this is an archetypal query, it must match all entities. + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + if Self::IS_ARCHETYPAL || !fetch.matches { + true + } else { + let ($($name,)*) = &mut fetch.fetch; + let ($($state,)*) = state; + // SAFETY: invariants are upheld by the caller. + false $(|| ($name.matches && unsafe { $name::matches(&$state, &mut $name.fetch, entity, table_row) }))* + } + } } #[expect( @@ -2845,7 +3139,6 @@ macro_rules! impl_anytuple_fetch { // SAFETY: defers to soundness of `$name: WorldQuery` impl unsafe impl<$($name: QueryData),*> QueryData for AnyOf<($($name,)*)> { const IS_READ_ONLY: bool = true $(&& $name::IS_READ_ONLY)*; - const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*; type ReadOnly = AnyOf<($($name::ReadOnly,)*)>; type Item<'w, 's> = ($(Option<$name::Item<'w, 's>>,)*); @@ -2862,24 +3155,46 @@ macro_rules! impl_anytuple_fetch { _fetch: &mut Self::Fetch<'w>, _entity: Entity, _table_row: TableRow + ) -> Self::Item<'w, 's> { + if !_fetch.matches { + ($(Option::<$name::Item<'w, 's>>::None,)*) + } else { + let ($($name,)*) = &mut _fetch.fetch; + let ($($state,)*) = _state; + ($( + // SAFETY: The invariants are required to be upheld by the caller. + ( $name.matches && unsafe { $name::matches(&$state, &mut $name.fetch, _entity, _table_row) }).then(|| + // SAFETY: the above guard verifies this subquery matches the given entity. + // other invariants are upheld by the caller. + unsafe { $name::fetch($state, &mut $name.fetch, _entity, _table_row) } + ), + )*) + } + } + + #[inline(always)] + unsafe fn try_fetch<'w, 's>( + _state: &'s Self::State, + _fetch: &mut Self::Fetch<'w>, + _entity: Entity, + _table_row: TableRow ) -> Option> { - let ($($name,)*) = _fetch; + let ($($name,)*) = &mut _fetch.fetch; let ($($state,)*) = _state; let result = ($( // SAFETY: The invariants are required to be upheld by the caller. - $name.1.then(|| unsafe { $name::fetch($state, &mut $name.0, _entity, _table_row) }).flatten(), + $name.matches.then(|| unsafe { $name::try_fetch($state, &mut $name.fetch, _entity, _table_row) }).flatten(), )*); // If this is an archetypal query, then it is guaranteed to return `Some`, // and we can help the compiler remove branches by checking the const `IS_ARCHETYPAL` first. - (Self::IS_ARCHETYPAL + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + (Self::IS_ARCHETYPAL || !_fetch.matches // We want to return `Some` if the query matches this entity, // which happens if at least one subquery returns `Some`. // So, fetch everything as usual, but if all the subqueries return `None` then return `None` instead. - || !matches!(result, ($(Option::>::None,)*)) - // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. - // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, - // so we treat them as matching for non-archetypal queries, as well. - || !(false $(|| $name.1)*)) + || !matches!(result, ($(Option::>::None,)*))) .then_some(result) } @@ -2956,6 +3271,7 @@ unsafe impl WorldQuery for NopWorldQuery { } const IS_DENSE: bool = D::IS_DENSE; + const IS_ARCHETYPAL: bool = true; #[inline(always)] unsafe fn set_archetype( @@ -2985,12 +3301,21 @@ unsafe impl WorldQuery for NopWorldQuery { ) -> bool { D::matches_component_set(state, set_contains_id) } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self::ReadOnly` is `Self` unsafe impl QueryData for NopWorldQuery { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = (); @@ -3005,8 +3330,7 @@ unsafe impl QueryData for NopWorldQuery { _fetch: &mut Self::Fetch<'w>, _entity: Entity, _table_row: TableRow, - ) -> Option> { - Some(()) + ) -> Self::Item<'w, 's> { } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -3045,6 +3369,7 @@ unsafe impl WorldQuery for PhantomData { // `PhantomData` does not match any components, so all components it matches // are stored in a Table (vacuous truth). const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; unsafe fn set_archetype<'w, 's>( _fetch: &mut Self::Fetch<'w>, @@ -3075,12 +3400,21 @@ unsafe impl WorldQuery for PhantomData { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self::ReadOnly` is `Self` unsafe impl QueryData for PhantomData { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = (); @@ -3094,8 +3428,7 @@ unsafe impl QueryData for PhantomData { _fetch: &mut Self::Fetch<'w>, _entity: Entity, _table_row: TableRow, - ) -> Option> { - Some(()) + ) -> Self::Item<'w, 's> { } fn iter_access(_state: &Self::State) -> impl Iterator> { @@ -3253,6 +3586,7 @@ mod tests { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -3285,13 +3619,22 @@ mod tests { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for NonReleaseQueryData { type ReadOnly = Self; const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type Item<'w, 's> = (); @@ -3306,8 +3649,7 @@ mod tests { _fetch: &mut Self::Fetch<'w>, _entity: Entity, _table_row: TableRow, - ) -> Option> { - Some(()) + ) -> Self::Item<'w, 's> { } fn iter_access(_state: &Self::State) -> impl Iterator> { diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index f8578d5d21705..e08ce19beefd1 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -1,15 +1,15 @@ use crate::{ - archetype::Archetype, + archetype::{Archetype, ArchetypeEntity}, change_detection::Tick, component::{Component, ComponentId, Components, StorageType}, entity::{Entities, Entity}, - query::{DebugCheckedUnwrap, FilteredAccess, StorageSwitch, WorldQuery}, + query::{DebugCheckedUnwrap, FilteredAccess, RangeExt, StorageSwitch, WorldQuery}, storage::{ComponentSparseSet, Table, TableRow}, world::{unsafe_world_cell::UnsafeWorldCell, World}, }; use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref}; use bevy_utils::prelude::DebugName; -use core::{cell::UnsafeCell, marker::PhantomData}; +use core::{cell::UnsafeCell, marker::PhantomData, ops::Range}; use variadics_please::all_tuples; /// Types that filter the results of a [`Query`]. @@ -71,46 +71,12 @@ use variadics_please::all_tuples; /// ``` /// /// [`Query`]: crate::system::Query -/// -/// # Safety -/// -/// The [`WorldQuery`] implementation must not take any mutable access. -/// This is the same safety requirement as [`ReadOnlyQueryData`](crate::query::ReadOnlyQueryData). #[diagnostic::on_unimplemented( message = "`{Self}` is not a valid `Query` filter", label = "invalid `Query` filter", note = "a `QueryFilter` typically uses a combination of `With` and `Without` statements" )] -pub unsafe trait QueryFilter: WorldQuery { - /// Returns true if (and only if) this Filter relies strictly on archetypes to limit which - /// components are accessed by the Query. - /// - /// This enables optimizations for [`QueryIter`](`crate::query::QueryIter`) that rely on knowing exactly how - /// many elements are being iterated (such as `Iterator::collect()`). - /// - /// If this is `true`, then [`QueryFilter::filter_fetch`] must always return true. - const IS_ARCHETYPAL: bool; - - /// Returns true if the provided [`Entity`] and [`TableRow`] should be included in the query results. - /// If false, the entity will be skipped. - /// - /// Note that this is called after already restricting the matched [`Table`]s and [`Archetype`]s to the - /// ones that are compatible with the Filter's access. - /// - /// Implementors of this method will generally either have a trivial `true` body (required for archetypal filters), - /// or access the necessary data within this function to make the final decision on filter inclusion. - /// - /// # 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 filter_fetch( - state: &Self::State, - fetch: &mut Self::Fetch<'_>, - entity: Entity, - table_row: TableRow, - ) -> bool; -} +pub trait QueryFilter: WorldQuery {} /// Filter that selects entities with a component `T`. /// @@ -143,7 +109,7 @@ pub struct With(PhantomData); /// SAFETY: /// `update_component_access` does not add any accesses. -/// This is sound because [`QueryFilter::filter_fetch`] does not access any components. +/// This is sound because [`WorldQuery::matches`] does not access any components. /// `update_component_access` adds a `With` filter for `T`. /// This is sound because `matches_component_set` returns whether the set contains the component. unsafe impl WorldQuery for With { @@ -167,6 +133,7 @@ unsafe impl WorldQuery for With { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype( @@ -199,14 +166,9 @@ unsafe impl WorldQuery for With { ) -> bool { set_contains_id(id) } -} -// SAFETY: WorldQuery impl performs no access at all -unsafe impl QueryFilter for With { - const IS_ARCHETYPAL: bool = true; - - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn matches( _state: &Self::State, _fetch: &mut Self::Fetch<'_>, _entity: Entity, @@ -216,6 +178,8 @@ unsafe impl QueryFilter for With { } } +impl QueryFilter for With {} + /// Filter that selects entities without a component `T`. /// /// This is the negation of [`With`]. @@ -244,7 +208,7 @@ pub struct Without(PhantomData); /// SAFETY: /// `update_component_access` does not add any accesses. -/// This is sound because [`QueryFilter::filter_fetch`] does not access any components. +/// This is sound because [`WorldQuery::matches`] does not access any components. /// `update_component_access` adds a `Without` filter for `T`. /// This is sound because `matches_component_set` returns whether the set does not contain the component. unsafe impl WorldQuery for Without { @@ -268,6 +232,7 @@ unsafe impl WorldQuery for Without { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype( @@ -300,14 +265,9 @@ unsafe impl WorldQuery for Without { ) -> bool { !set_contains_id(id) } -} -// SAFETY: WorldQuery impl performs no access at all -unsafe impl QueryFilter for Without { - const IS_ARCHETYPAL: bool = true; - - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn matches( _state: &Self::State, _fetch: &mut Self::Fetch<'_>, _entity: Entity, @@ -317,6 +277,8 @@ unsafe impl QueryFilter for Without { } } +impl QueryFilter for Without {} + /// A filter that tests if any of the given filters apply. /// /// This is useful for example if a system with multiple components in a query only wants to run @@ -350,20 +312,12 @@ unsafe impl QueryFilter for Without { pub struct Or(PhantomData); #[doc(hidden)] -pub struct OrFetch<'w, T: WorldQuery> { - fetch: T::Fetch<'w>, +#[derive(Clone)] +pub struct OrFetch { + fetch: T, matches: bool, } -impl Clone for OrFetch<'_, T> { - fn clone(&self) -> Self { - Self { - fetch: self.fetch.clone(), - matches: self.matches, - } - } -} - macro_rules! impl_or_query_filter { ($(#[$meta:meta])* $(($filter: ident, $state: ident)),*) => { $(#[$meta])* @@ -383,35 +337,50 @@ macro_rules! impl_or_query_filter { clippy::unused_unit, reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." )] + #[allow( + unused_mut, + reason = "Zero-length tuples won't mutate any of the parameters." + )] /// SAFETY: - /// [`QueryFilter::filter_fetch`] accesses are a subset of the subqueries' accesses + /// [`WorldQuery::matches`], [`WorldQuery::find_table_chunk`], and [`WorldQuery::find_archetype_chunk`] only + /// call their respective methods on subqueries, so they can't perform mutable access unless subqueries are invalid. + /// [`WorldQuery::matches`] accesses are a subset of the subqueries' accesses /// This is sound because `update_component_access` adds accesses according to the implementations of all the subqueries. /// `update_component_access` replace the filters with a disjunction where every element is a conjunction of the previous filters and the filters of one of the subqueries. /// This is sound because `matches_component_set` returns a disjunction of the results of the subqueries' implementations. unsafe impl<$($filter: QueryFilter),*> WorldQuery for Or<($($filter,)*)> { - type Fetch<'w> = ($(OrFetch<'w, $filter>,)*); + type Fetch<'w> = OrFetch<($(OrFetch<$filter::Fetch<'w>>,)*)>; type State = ($($filter::State,)*); fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { - let ($($filter,)*) = fetch; - ($( - OrFetch { - fetch: $filter::shrink_fetch($filter.fetch), - matches: $filter.matches - }, - )*) + let ($($filter,)*) = fetch.fetch; + OrFetch { + fetch: + ($( + OrFetch { + fetch: $filter::shrink_fetch($filter.fetch), + matches: $filter.matches + }, + )*), + matches: fetch.matches, + } } const IS_DENSE: bool = true $(&& $filter::IS_DENSE)*; + const IS_ARCHETYPAL: bool = true $(&& $filter::IS_ARCHETYPAL)*; #[inline] unsafe fn init_fetch<'w, 's>(world: UnsafeWorldCell<'w>, state: &'s Self::State, last_run: Tick, this_run: Tick) -> Self::Fetch<'w> { let ($($filter,)*) = state; - ($(OrFetch { - // SAFETY: The invariants are upheld by the caller. - fetch: unsafe { $filter::init_fetch(world, $filter, last_run, this_run) }, - matches: false, - },)*) + OrFetch { + fetch: + ($(OrFetch { + // SAFETY: The invariants are upheld by the caller. + fetch: unsafe { $filter::init_fetch(world, $filter, last_run, this_run) }, + matches: false, + },)*), + matches: false + } } #[inline] @@ -421,10 +390,12 @@ macro_rules! impl_or_query_filter { if Self::IS_ARCHETYPAL { return; } - let ($($filter,)*) = fetch; + let ($($filter,)*) = &mut fetch.fetch; let ($($state,)*) = state; + fetch.matches = false; $( $filter.matches = $filter::matches_component_set($state, &|id| table.has_column(id)); + fetch.matches = fetch.matches || $filter.matches; if $filter.matches { // SAFETY: The invariants are upheld by the caller. unsafe { $filter::set_table(&mut $filter.fetch, $state, table); } @@ -444,10 +415,12 @@ macro_rules! impl_or_query_filter { if Self::IS_ARCHETYPAL { return; } - let ($($filter,)*) = fetch; + let ($($filter,)*) = &mut fetch.fetch; let ($($state,)*) = &state; + fetch.matches = false; $( $filter.matches = $filter::matches_component_set($state, &|id| archetype.contains(id)); + fetch.matches = fetch.matches || $filter.matches; if $filter.matches { // SAFETY: The invariants are upheld by the caller. unsafe { $filter::set_archetype(&mut $filter.fetch, $state, archetype, table); } @@ -489,81 +462,92 @@ macro_rules! impl_or_query_filter { let ($($filter,)*) = state; false $(|| $filter::matches_component_set($filter, set_contains_id))* } - } - - #[expect( - clippy::allow_attributes, - reason = "This is a tuple-related macro; as such the lints below may not always apply." - )] - #[allow( - non_snake_case, - reason = "The names of some variables are provided by the macro's caller, not by us." - )] - #[allow( - unused_variables, - reason = "Zero-length tuples won't use any of the parameters." - )] - $(#[$meta])* - // SAFETY: This only performs access that subqueries perform, and they impl `QueryFilter` and so perform no mutable access. - unsafe impl<$($filter: QueryFilter),*> QueryFilter for Or<($($filter,)*)> { - const IS_ARCHETYPAL: bool = true $(&& $filter::IS_ARCHETYPAL)*; - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn find_table_chunk( state: &Self::State, fetch: &mut Self::Fetch<'_>, - entity: Entity, - table_row: TableRow - ) -> bool { - let ($($state,)*) = state; - let ($($filter,)*) = fetch; - // If this is an archetypal query, then it is guaranteed to return true, - // and we can help the compiler remove branches by checking the const `IS_ARCHETYPAL` first. - (Self::IS_ARCHETYPAL - // SAFETY: The invariants are upheld by the caller. - $(|| ($filter.matches && unsafe { $filter::filter_fetch($state, &mut $filter.fetch, entity, table_row) }))* - // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. - // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, - // so we treat them as matching for non-archetypal queries, as well. - || !(false $(|| $filter.matches)*)) + table_entities: &[Entity], + rows: Range, + ) -> Range { + // If this is an archetypal query, it must match all entities. + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + if Self::IS_ARCHETYPAL || !fetch.matches { + rows + } else { + let ($($filter,)*) = &mut fetch.fetch; + let ($($state,)*) = state; + let mut new_rows = rows.end..rows.end; + $( + if $filter.matches { + // SAFETY: invariants are upheld by the caller. + new_rows = new_rows.union_or_first(unsafe { $filter::find_table_chunk($state, &mut $filter.fetch, table_entities, rows.clone()) }); + } + )* + new_rows + } } - } - }; -} -macro_rules! impl_tuple_query_filter { - ($(#[$meta:meta])* $(($name: ident, $state: ident)),*) => { - #[expect( - clippy::allow_attributes, - reason = "This is a tuple-related macro; as such the lints below may not always apply." - )] - #[allow( - non_snake_case, - reason = "The names of some variables are provided by the macro's caller, not by us." - )] - #[allow( - unused_variables, - reason = "Zero-length tuples won't use any of the parameters." - )] - $(#[$meta])* - // SAFETY: This only performs access that subqueries perform, and they impl `QueryFilter` and so perform no mutable access. - unsafe impl<$($name: QueryFilter),*> QueryFilter for ($($name,)*) { - const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*; + #[inline] + unsafe fn find_archetype_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + archetype_entities: &[ArchetypeEntity], + mut indices: Range, + ) -> Range { + // If this is an archetypal query, it must match all entities. + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + if Self::IS_ARCHETYPAL || !fetch.matches { + indices + } else { + let ($($filter,)*) = &mut fetch.fetch; + let ($($state,)*) = state; + let mut new_indices = indices.end..indices.end; + $( + if $filter.matches { + // SAFETY: invariants are upheld by the caller. + new_indices = new_indices.union_or_first(unsafe { $filter::find_archetype_chunk($state, &mut $filter.fetch, archetype_entities, indices.clone()) }); + } + )* + new_indices + } + } - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn matches( state: &Self::State, fetch: &mut Self::Fetch<'_>, entity: Entity, - table_row: TableRow + table_row: TableRow, ) -> bool { - let ($($state,)*) = state; - let ($($name,)*) = fetch; - // SAFETY: The invariants are upheld by the caller. - true $(&& unsafe { $name::filter_fetch($state, $name, entity, table_row) })* + // If this is an archetypal query, it must match all entities. + // If *none* of the subqueries matched the archetype, then this archetype was added in a transmute. + // We must treat those as matching in order to be consistent with `size_hint` for archetypal queries, + // so we treat them as matching for non-archetypal queries, as well. + if Self::IS_ARCHETYPAL || !fetch.matches { + true + } else { + let ($($filter,)*) = &mut fetch.fetch; + let ($($state,)*) = state; + // SAFETY: invariants are upheld by the caller. + false $(|| ($filter.matches && unsafe { $filter::matches(&$state, &mut $filter.fetch, entity, table_row) }))* + } } } + $(#[$meta])* + impl<$($filter: QueryFilter),*> QueryFilter for Or<($($filter,)*)> {} + }; +} + +macro_rules! impl_tuple_query_filter { + ($(#[$meta:meta])* $($name: ident),*) => { + $(#[$meta])* + impl<$($name: QueryFilter),*> QueryFilter for ($($name,)*) {} }; } @@ -572,8 +556,7 @@ all_tuples!( impl_tuple_query_filter, 0, 15, - F, - S + F ); all_tuples!( #[doc(fake_variadic)] @@ -591,7 +574,7 @@ pub struct Allow(PhantomData); /// SAFETY: /// `update_component_access` does not add any accesses. -/// This is sound because [`QueryFilter::filter_fetch`] does not access any components. +/// This is sound because [`WorldQuery::matches`] does not access any components. /// `update_component_access` adds an archetypal filter for `T`. /// This is sound because it doesn't affect the query unsafe impl WorldQuery for Allow { @@ -605,6 +588,7 @@ unsafe impl WorldQuery for Allow { // Even if the component is sparse, this implementation doesn't do anything with it const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype(_: &mut (), _: &ComponentId, _: &Archetype, _: &Table) {} @@ -629,23 +613,20 @@ unsafe impl WorldQuery for Allow { // Allow always matches true } -} - -// SAFETY: WorldQuery impl performs no access at all -unsafe impl QueryFilter for Allow { - const IS_ARCHETYPAL: bool = true; - #[inline(always)] - unsafe fn filter_fetch( - _: &Self::State, - _: &mut Self::Fetch<'_>, - _: Entity, - _: TableRow, + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, ) -> bool { true } } +impl QueryFilter for Allow {} + /// A filter on a component that only retains results the first time after they have been added. /// /// A common use for this filter is one-time initialization. @@ -737,7 +718,7 @@ impl Clone for AddedFetch<'_, T> { } /// SAFETY: -/// [`QueryFilter::filter_fetch`] accesses a single component in a readonly way. +/// [`WorldQuery::matches`] accesses a single component in a readonly way. /// This is sound because `update_component_access` adds read access for that component and panics when appropriate. /// `update_component_access` adds a `With` filter for a component. /// This is sound because `matches_component_set` returns whether the set contains that component. @@ -778,6 +759,7 @@ unsafe impl WorldQuery for Added { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = false; #[inline] unsafe fn set_archetype<'w, 's>( @@ -832,13 +814,9 @@ unsafe impl WorldQuery for Added { ) -> bool { set_contains_id(id) } -} -// SAFETY: WorldQuery impl performs only read access on ticks -unsafe impl QueryFilter for Added { - const IS_ARCHETYPAL: bool = false; - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn matches( _state: &Self::State, fetch: &mut Self::Fetch<'_>, entity: Entity, @@ -869,6 +847,8 @@ unsafe impl QueryFilter for Added { } } +impl QueryFilter for Added {} + /// A filter on a component that only retains results the first time after they have been added or mutably dereferenced. /// /// A common use for this filter is avoiding redundant work when values have not changed. @@ -964,7 +944,7 @@ impl Clone for ChangedFetch<'_, T> { } /// SAFETY: -/// `fetch` accesses a single component in a readonly way. +/// `matches` accesses a single component in a readonly way. /// This is sound because `update_component_access` add read access for that component and panics when appropriate. /// `update_component_access` adds a `With` filter for a component. /// This is sound because `matches_component_set` returns whether the set contains that component. @@ -1005,6 +985,7 @@ unsafe impl WorldQuery for Changed { StorageType::SparseSet => false, } }; + const IS_ARCHETYPAL: bool = false; #[inline] unsafe fn set_archetype<'w, 's>( @@ -1059,14 +1040,9 @@ unsafe impl WorldQuery for Changed { ) -> bool { set_contains_id(id) } -} -// SAFETY: WorldQuery impl performs only read access on ticks -unsafe impl QueryFilter for Changed { - const IS_ARCHETYPAL: bool = false; - - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn matches( _state: &Self::State, fetch: &mut Self::Fetch<'_>, entity: Entity, @@ -1097,6 +1073,8 @@ unsafe impl QueryFilter for Changed { } } +impl QueryFilter for Changed {} + /// A filter that only retains results the first time after the entity has been spawned. /// /// A common use for this filter is one-time initialization. @@ -1126,7 +1104,7 @@ unsafe impl QueryFilter for Changed { /// # use bevy_ecs::query::SpawnDetails; /// /// fn system1(query: Query) { -/// for entity in &query { /* entity spawned */ } +/// for entity in &query { /* entity spawned*/ } /// } /// /// fn system2(query: Query<(Entity, SpawnDetails)>) { @@ -1189,6 +1167,7 @@ unsafe impl WorldQuery for Spawned { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = false; #[inline] unsafe fn set_archetype<'w, 's>( @@ -1214,14 +1193,9 @@ unsafe impl WorldQuery for Spawned { fn matches_component_set(_state: &(), _set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { true } -} -// SAFETY: WorldQuery impl accesses no components or component ticks -unsafe impl QueryFilter for Spawned { - const IS_ARCHETYPAL: bool = false; - - #[inline(always)] - unsafe fn filter_fetch( + #[inline] + unsafe fn matches( _state: &Self::State, fetch: &mut Self::Fetch<'_>, entity: Entity, @@ -1238,12 +1212,14 @@ unsafe impl QueryFilter for Spawned { } } +impl QueryFilter for Spawned {} + /// A marker trait to indicate that the filter works at an archetype level. /// /// This is needed to implement [`ExactSizeIterator`] for /// [`QueryIter`](crate::query::QueryIter) that contains archetype-level filters. /// -/// The trait must only be implemented for filters where its corresponding [`QueryFilter::IS_ARCHETYPAL`] +/// The trait must only be implemented for filters where its corresponding [`WorldQuery::IS_ARCHETYPAL`] /// is [`prim@true`]. As such, only the [`With`] and [`Without`] filters can implement the trait. /// [Tuples](prim@tuple) and [`Or`] filters are automatically implemented with the trait only if its containing types /// also implement the same trait. diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 0bd178ba03727..2f49db570cce3 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -216,38 +216,60 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { table, ); - let entities = table.entities(); - for row in rows { - // SAFETY: Caller assures `row` in range of the current archetype. - let entity = unsafe { entities.get_unchecked(row as usize) }; - // SAFETY: This is from an exclusive range, so it can't be max. - let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(row)) }; - - // SAFETY: set_table was called prior. - // Caller assures `row` in range of the current archetype. - let fetched = unsafe { - !F::filter_fetch( + let table_entities = table.entities(); + + let mut next_chunk = rows.clone(); + while next_chunk.start < rows.end { + // SAFETY: + // - `set_table` was called prior + // - since `find_table_chunk` can only narrow ranges, and the caller + // assures `rows` is in range of the table, `next_chunk` must be as well. + next_chunk = unsafe { + D::find_table_chunk( + &self.query_state.fetch_state, + &mut self.cursor.fetch, + table_entities, + next_chunk, + ) + }; + // SAFETY: + // - `set_table` was called prior + // - since `find_table_chunk` can only narrow ranges, and the caller + // assures `rows` is in range of the table, `next_chunk` must be as well. + next_chunk = unsafe { + F::find_table_chunk( &self.query_state.filter_state, &mut self.cursor.filter, - *entity, - row, + table_entities, + next_chunk, ) }; - if fetched { - continue; + + if next_chunk.is_empty() { + break; } - // SAFETY: set_table was called prior. - // Caller assures `row` in range of the current archetype. - if let Some(item) = D::fetch( - &self.query_state.fetch_state, - &mut self.cursor.fetch, - *entity, - row, - ) { + for row in next_chunk.clone() { + // SAFETY: Caller assures `row` in range of the current archetype. + let entity = unsafe { table_entities.get_unchecked(row as usize) }; + // SAFETY: This is from an exclusive range, so it can't be max. + let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(row)) }; + + // SAFETY: + // - set_table was called prior. + // - find_table_chunk was called prior, and this row was in the returned range. + let item = D::fetch( + &self.query_state.fetch_state, + &mut self.cursor.fetch, + *entity, + row, + ); accum = func(accum, item); } + + next_chunk = next_chunk.end..rows.end; } + accum } @@ -286,37 +308,59 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { table, ); - let entities = archetype.entities(); - for index in indices { - // SAFETY: Caller assures `index` in range of the current archetype. - let archetype_entity = unsafe { entities.get_unchecked(index as usize) }; + let archetype_entities = archetype.entities(); - // SAFETY: set_archetype was called prior. - // Caller assures `index` in range of the current archetype. - let fetched = unsafe { - !F::filter_fetch( + let mut next_chunk = indices.clone(); + while next_chunk.start < indices.end { + // SAFETY: + // - `set_archetype` was called prior + // - since `find_archetype_chunk` can only narrow ranges, and the caller + // assures `indices` is in range of the archetype, `next_chunk` must be as well. + next_chunk = unsafe { + D::find_archetype_chunk( + &self.query_state.fetch_state, + &mut self.cursor.fetch, + archetype_entities, + next_chunk, + ) + }; + // SAFETY: + // - `set_archetype` was called prior + // - since `find_archetype_chunk` can only narrow ranges, and the caller + // assures `indices` is in range of the archetype, `next_chunk` must be as well. + next_chunk = unsafe { + F::find_archetype_chunk( &self.query_state.filter_state, &mut self.cursor.filter, - archetype_entity.id(), - archetype_entity.table_row(), + archetype_entities, + next_chunk, ) }; - if fetched { - continue; + + if next_chunk.is_empty() { + break; } - // SAFETY: set_archetype was called prior, `index` is an archetype index in range of the current archetype - // Caller assures `index` in range of the current archetype. - if let Some(item) = unsafe { - D::fetch( - &self.query_state.fetch_state, - &mut self.cursor.fetch, - archetype_entity.id(), - archetype_entity.table_row(), - ) - } { + for index in next_chunk.clone() { + // SAFETY: by induction (see above) index must be in range of the archetype. + let archetype_entity = unsafe { archetype_entities.get_unchecked(index as usize) }; + + // SAFETY: + // - `set_archetype` was called prior. + // - `find_archetype_chunk` was called prior, and `index` is in the returned range. + let item = unsafe { + D::fetch( + &self.query_state.fetch_state, + &mut self.cursor.fetch, + archetype_entity.id(), + archetype_entity.table_row(), + ) + }; + accum = func(accum, item); } + + next_chunk = next_chunk.end..indices.end; } accum } @@ -361,37 +405,62 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { archetype, table, ); - let entities = table.entities(); - for row in rows { - // SAFETY: Caller assures `row` in range of the current archetype. - let entity = unsafe { *entities.get_unchecked(row as usize) }; - // SAFETY: This is from an exclusive range, so it can't be max. - let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(row)) }; - - // SAFETY: set_table was called prior. - // Caller assures `row` in range of the current archetype. - let filter_matched = unsafe { - F::filter_fetch( + + let table_entities = table.entities(); + + let mut next_chunk = rows.clone(); + while next_chunk.start < rows.end { + // SAFETY: + // - `set_table` was called prior + // - since `find_table_chunk` can only narrow ranges, and the caller + // assures `rows` is in range of the archetype, `next_chunk` must be as well. + next_chunk = unsafe { + D::find_table_chunk( + &self.query_state.fetch_state, + &mut self.cursor.fetch, + table_entities, + next_chunk, + ) + }; + // SAFETY: + // - `set_table` was called prior + // - since `find_table_chunk` can only narrow ranges, and the caller + // assures `rows` is in range of the archetype, `next_chunk` must be as well. + next_chunk = unsafe { + F::find_table_chunk( &self.query_state.filter_state, &mut self.cursor.filter, - entity, - row, + table_entities, + next_chunk, ) }; - if !filter_matched { - continue; + + if next_chunk.is_empty() { + break; } - // SAFETY: set_table was called prior. - // Caller assures `row` in range of the current archetype. - if let Some(item) = D::fetch( - &self.query_state.fetch_state, - &mut self.cursor.fetch, - entity, - row, - ) { + for row in next_chunk.clone() { + // SAFETY: by induction (see above) `row` must be in range of the table. + let entity = unsafe { *table_entities.get_unchecked(row as usize) }; + // SAFETY: This is from an exclusive range, so it can't be max. + let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(row)) }; + + // SAFETY: + // - `set_archetype` was called prior. + // - `find_table_chunk` was called prior, and `row` is in the returned range. + let item = unsafe { + D::fetch( + &self.query_state.fetch_state, + &mut self.cursor.fetch, + entity, + row, + ) + }; + accum = func(accum, item); } + + next_chunk = next_chunk.end..rows.end; } accum } @@ -1067,16 +1136,14 @@ where // The entity list has already been filtered by the query lens, so we forego filtering here. // SAFETY: - // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype + // - set_archetype was called prior, `location.table_row` is an archetype index in range of the current archetype // - fetch is only called once for each entity. - unsafe { - D::fetch( - &self.query_state.fetch_state, - &mut self.fetch, - entity, - location.table_row, - ) - } + D::try_fetch( + &self.query_state.fetch_state, + &mut self.fetch, + entity, + location.table_row, + ) } } @@ -1246,24 +1313,25 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> // SAFETY: set_archetype was called prior. // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d - if unsafe { - F::filter_fetch( + let matches_filter = unsafe { + F::matches( &query_state.filter_state, filter, entity, location.table_row, ) - } { - // SAFETY: - // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype - // - fetch is only called once for each entity. - let item = unsafe { - D::fetch(&query_state.fetch_state, fetch, entity, location.table_row) - }; - if let Some(item) = item { - return Some(item); - } - } + }; + + return matches_filter + .then(|| { + // SAFETY: + // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype + // - fetch is only called once for each entity. + unsafe { + D::try_fetch(&query_state.fetch_state, fetch, entity, location.table_row) + } + }) + .flatten(); } None } @@ -2042,12 +2110,12 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype // - fetch is only called once for each entity. unsafe { - D::fetch( + Some(D::fetch( &self.query_state.fetch_state, &mut self.fetch, entity, location.table_row, - ) + )) } } @@ -2480,34 +2548,52 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { let index = self.current_row - 1; if self.is_dense { // SAFETY: This must have been called previously in `next` as `current_row > 0` - let entity = unsafe { self.table_entities.get_unchecked(index as usize) }; - // SAFETY: - // - `set_table` must have been called previously either in `next` or before it. - // - `*entity` and `index` are in the current table. - unsafe { - D::fetch( - &query_state.fetch_state, - &mut self.fetch, - *entity, - // SAFETY: This is from an exclusive range, so it can't be max. - TableRow::new(NonMaxU32::new_unchecked(index)), - ) - } + let entity = unsafe { *self.table_entities.get_unchecked(index as usize) }; + // SAFETY: This is from an exclusive range, so it can't be max. + let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(index)) }; + + // SAFETY: set_table was called prior. + // `row` is a table row in range of the current table, + let matches_filter = + unsafe { F::matches(&query_state.filter_state, &mut self.filter, entity, row) }; + + matches_filter + // SAFETY: + // - `set_table` must have been called previously either in `next` or before it. + // - `*entity` and `index` are in the current table. + .then(|| unsafe { + D::try_fetch(&query_state.fetch_state, &mut self.fetch, entity, row) + }) + .flatten() } else { // SAFETY: This must have been called previously in `next` as `current_row > 0` let archetype_entity = unsafe { self.archetype_entities.get_unchecked(index as usize) }; - // SAFETY: - // - `set_archetype` must have been called previously either in `next` or before it. - // - `archetype_entity.id()` and `archetype_entity.table_row()` are in the current archetype. - unsafe { - D::fetch( - &query_state.fetch_state, - &mut self.fetch, + + // SAFETY: set_archetype was called prior. + // `archetype_entity.table_row()` is an archetype index row in range of the current archetype, + let matches_filter = unsafe { + F::matches( + &query_state.filter_state, + &mut self.filter, archetype_entity.id(), archetype_entity.table_row(), ) - } + }; + + matches_filter + // SAFETY: + // - `set_archetype` must have been called previously either in `next` or before it. + // - `archetype_entity.id()` and `archetype_entity.table_row()` are in the current archetype. + .then(|| unsafe { + D::try_fetch( + &query_state.fetch_state, + &mut self.fetch, + archetype_entity.id(), + archetype_entity.table_row(), + ) + }) + .flatten() } } else { None @@ -2568,12 +2654,17 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // SAFETY: set_table was called prior. // `current_row` is a table row in range of the current table, because if it was not, then the above would have been executed. let entity = - unsafe { self.table_entities.get_unchecked(self.current_row as usize) }; + unsafe { *self.table_entities.get_unchecked(self.current_row as usize) }; // SAFETY: The row is less than the u32 len, so it must not be max. let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(self.current_row)) }; self.current_row += 1; - if !F::filter_fetch(&query_state.filter_state, &mut self.filter, *entity, row) { + // SAFETY: set_table was called prior. + // `row` is a table row in range of the current table, + let matches_filter = + unsafe { F::matches(&query_state.filter_state, &mut self.filter, entity, row) }; + + if !matches_filter { continue; } @@ -2583,7 +2674,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // because if it was not, then the above would have been executed. // - fetch is only called once for each `entity`. let item = - unsafe { D::fetch(&query_state.fetch_state, &mut self.fetch, *entity, row) }; + unsafe { D::try_fetch(&query_state.fetch_state, &mut self.fetch, entity, row) }; + if let Some(item) = item { return Some(item); } @@ -2626,12 +2718,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { }; self.current_row += 1; - if !F::filter_fetch( - &query_state.filter_state, - &mut self.filter, - archetype_entity.id(), - archetype_entity.table_row(), - ) { + // SAFETY: set_archetype was called prior. + // `archetype_entity.table_row()` is an archetype index row in range of the current archetype, + let matches_filter = unsafe { + F::matches( + &query_state.filter_state, + &mut self.filter, + archetype_entity.id(), + archetype_entity.table_row(), + ) + }; + + if !matches_filter { continue; } @@ -2641,13 +2739,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // because if it was not, then the if above would have been executed. // - fetch is only called once for each `archetype_entity`. let item = unsafe { - D::fetch( + D::try_fetch( &query_state.fetch_state, &mut self.fetch, archetype_entity.id(), archetype_entity.table_row(), ) }; + if let Some(item) = item { return Some(item); } diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 065eadd24db88..b55a91db9d1c2 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -841,6 +841,7 @@ mod tests { } const IS_DENSE: bool = true; + const IS_ARCHETYPAL: bool = true; #[inline] unsafe fn set_archetype<'w, 's>( @@ -881,12 +882,21 @@ mod tests { ) -> bool { true } + + #[inline] + unsafe fn matches( + _state: &Self::State, + _fetch: &mut Self::Fetch<'_>, + _entity: Entity, + _table_row: TableRow, + ) -> bool { + true + } } /// SAFETY: `Self` is the same as `Self::ReadOnly` unsafe impl QueryData for ReadsRData { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = true; type ReadOnly = Self; type Item<'w, 's> = (); @@ -901,8 +911,7 @@ mod tests { _fetch: &mut Self::Fetch<'w>, _entity: Entity, _table_row: TableRow, - ) -> Option> { - Some(()) + ) -> Self::Item<'w, 's> { } fn iter_access( diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index fef5b0257f167..0bd3f2e2bf687 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/crates/bevy_ecs/src/query/world_query.rs @@ -1,11 +1,15 @@ +use core::{mem, ops::Range}; + use crate::{ - archetype::Archetype, + archetype::{Archetype, ArchetypeEntity}, change_detection::Tick, component::{ComponentId, Components}, + entity::Entity, query::FilteredAccess, - storage::Table, + storage::{Table, TableRow}, world::{unsafe_world_cell::UnsafeWorldCell, World}, }; +use nonmax::NonMaxU32; use variadics_please::all_tuples; /// Types that can be used as parameters in a [`Query`]. @@ -14,11 +18,11 @@ use variadics_please::all_tuples; /// # Safety /// /// Implementor must ensure that -/// [`update_component_access`], [`QueryData::provide_extra_access`], [`matches_component_set`], [`QueryData::fetch`], [`QueryFilter::filter_fetch`] and [`init_fetch`] +/// [`update_component_access`], [`QueryData::provide_extra_access`], [`matches_component_set`], [`QueryData::fetch`], [`matches`] and [`init_fetch`] /// obey the following: /// /// - For each component mutably accessed by [`QueryData::fetch`], [`update_component_access`] or [`QueryData::provide_extra_access`] should add write access unless read or write access has already been added, in which case it should panic. -/// - For each component readonly accessed by [`QueryData::fetch`] or [`QueryFilter::filter_fetch`], [`update_component_access`] or [`QueryData::provide_extra_access`] should add read access unless write access has already been added, in which case it should panic. +/// - For each component readonly accessed by [`QueryData::fetch`] or [`matches`], [`update_component_access`] or [`QueryData::provide_extra_access`] should add read access unless write access has already been added, in which case it should panic. /// - If `fetch` mutably accesses the same component twice, [`update_component_access`] should panic. /// - [`update_component_access`] may not add a `Without` filter for a component unless [`matches_component_set`] always returns `false` when the component set contains that component. /// - [`update_component_access`] may not add a `With` filter for a component unless [`matches_component_set`] always returns `false` when the component set doesn't contain that component. @@ -32,15 +36,28 @@ use variadics_please::all_tuples; /// /// When implementing [`update_component_access`], note that `add_read` and `add_write` both also add a `With` filter, whereas `extend_access` does not change the filters. /// +/// If [`IS_ARCHETYPAL`] is `true`, [`matches`] must return `true` for all inputs. +/// +/// When implementing [`find_table_chunk`] and [`find_archetype_chunk`], their return values must be +/// _well formed_ subsets of `rows` and `indices` respectively, even if returning an empty range. +/// What this means is that given a range `(start..end)`, the returned range `(chunk_start..chunk_end)` +/// must satisfy `start <= chunk_start <= chunk_end <= end`. Additionally, they must match the behavior +/// of [`matches`], i.e. calling that method on any row/index in the returned range must return `true`. +/// +/// [`find_table_chunk`], [`find_archetype_chunk`], and [`matches`] must not mutably access any world data. +/// /// [`QueryData::provide_extra_access`]: crate::query::QueryData::provide_extra_access /// [`QueryData::fetch`]: crate::query::QueryData::fetch -/// [`QueryFilter::filter_fetch`]: crate::query::QueryFilter::filter_fetch /// [`init_fetch`]: Self::init_fetch /// [`matches_component_set`]: Self::matches_component_set /// [`Query`]: crate::system::Query /// [`update_component_access`]: Self::update_component_access /// [`QueryData`]: crate::query::QueryData /// [`QueryFilter`]: crate::query::QueryFilter +/// [`find_table_chunk`]: crate::query::WorldQuery::find_table_chunk +/// [`find_archetype_chunk`]: crate::query::WorldQuery::find_archetype_chunk +/// [`matches`]: crate::query::WorldQuery::matches +/// [`IS_ARCHETYPAL`]: crate::query::WorldQuery::IS_ARCHETYPAL pub unsafe trait WorldQuery { /// Per archetype/table state retrieved by this [`WorldQuery`] to compute [`Self::Item`](crate::query::QueryData::Item) for each entity. type Fetch<'w>: Clone; @@ -80,6 +97,15 @@ pub unsafe trait WorldQuery { /// iterators. const IS_DENSE: bool; + /// Returns true if (and only if) this query term relies strictly on archetypes to limit which + /// entities are accessed by the Query. + /// + /// This enables optimizations for [`QueryIter`](`crate::query::QueryIter`) that rely on knowing exactly how + /// many elements are being iterated (such as `Iterator::collect()`). + /// + /// If this is `true`, then [`WorldQuery::matches`] must always return `true`. + const IS_ARCHETYPAL: bool; + /// Adjusts internal state to account for the next [`Archetype`]. This will always be called on /// archetypes that match this [`WorldQuery`]. /// @@ -131,6 +157,176 @@ pub unsafe trait WorldQuery { state: &Self::State, set_contains_id: &impl Fn(ComponentId) -> bool, ) -> bool; + + /// Searches `rows` for the next contiguous range of indices + /// this `WorldQuery` matches. + /// + /// This can be used to speed up query iteration for custom [`WorldQuery`] + /// implementations, for example by using an acceleration structure + /// to seek ahead in the query for the next valid entity. + /// + /// The default implementation will simply defer to [`WorldQuery::matches`], + /// calling the method for each row until it finds the next whole chunk. + /// Custom implementations should match this behavior, i.e. calling [`WorldQuery::matches`] + /// on any row in the returned range should return `true`. + /// + /// # Safety + /// - Must always be called _after_ [`WorldQuery::set_table`]. + /// - `rows` must be in the range of the current table. + #[inline] + unsafe fn find_table_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + table_entities: &[Entity], + rows: Range, + ) -> Range { + if Self::IS_ARCHETYPAL { + rows + } else { + rows.find_chunk(|index| { + let table_row = TableRow::new(index); + // SAFETY: caller guarantees `indices` is in range of `table` + let entity = unsafe { *table_entities.get_unchecked(index.get() as usize) }; + // SAFETY: invariants upheld by caller + unsafe { Self::matches(state, fetch, entity, table_row) } + }) + } + } + + /// Searches `indices` for the next contiguous range of indices + /// this `WorldQuery` matches. + /// + /// This can be used to speed up query iteration for custom [`WorldQuery`] + /// implementations, for example by using an acceleration structure + /// to seek ahead in the query for the next valid entity. + /// + /// The default implementation will simply defer to [`WorldQuery::matches`], + /// calling the method for each row until it finds the next whole chunk. + /// Custom implementations should match this behavior, i.e. calling [`WorldQuery::matches`] + /// on any index in the returned range should return `true`. + /// + /// # Safety + /// - Must always be called _after_ [`WorldQuery::set_archetype`]. + /// - `indices` must be in the range of the current archetype. + #[inline] + unsafe fn find_archetype_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + archetype_entities: &[ArchetypeEntity], + indices: Range, + ) -> Range { + if Self::IS_ARCHETYPAL { + indices + } else { + indices.find_chunk(|index| { + // SAFETY: caller guarantees `indices` is in range of `archetype` + let archetype_entity = + unsafe { archetype_entities.get_unchecked(index.get() as usize) }; + + // SAFETY: invariants upheld by caller + unsafe { + Self::matches( + state, + fetch, + archetype_entity.id(), + archetype_entity.table_row(), + ) + } + }) + } + } + + /// Returns true if the provided [`Entity`] and [`TableRow`] should be included in the query results. + /// If false, the entity will be skipped. + /// + /// Note that this is called after already restricting the matched [`Table`]s and [`Archetype`]s to the + /// ones that are compatible with the Filter's access. + /// + /// Implementors of this method will generally either have a trivial `true` body (required for archetypal filters), + /// or access the necessary data within this function to make the final decision on filter inclusion. + /// + /// # 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 matches( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + entity: Entity, + table_row: TableRow, + ) -> bool; +} + +/// Extension trait for `Range`, adding convenience methods for use with [`WorldQuery::find_table_chunk`] +/// and [`WorldQuery::find_archetype_chunk`]. +pub trait RangeExt { + /// Returns the union of `self` and `other` if they overlap, otherwise returns + /// whichever is first. + fn union_or_first(self, other: Self) -> Self; + /// Returns the next contiguous segment for which `func` returns true for all + /// indices in the segment. + fn find_chunk bool>(self, func: F) -> Self; +} + +impl RangeExt for Range { + #[inline] + fn union_or_first(mut self, mut other: Self) -> Self { + if self.start > other.start { + mem::swap(&mut self, &mut other); + } + + if self.end >= other.start && self.end < other.end { + self.end = other.end; + } + + self + } + + #[inline] + fn find_chunk bool>(mut self, mut func: F) -> Self { + let mut index = self.start; + while index < self.end { + // SAFETY: index is taken from an exclusive range, so it can't be max + let nonmax_index = unsafe { NonMaxU32::new_unchecked(index) }; + let matches = func(nonmax_index); + if matches { + break; + } + index = index.wrapping_add(1); + } + self.start = index; + + while index < self.end { + // SAFETY: index is taken from an exclusive range, so it can't be max + let nonmax_index = unsafe { NonMaxU32::new_unchecked(index) }; + let matches = func(nonmax_index); + if !matches { + break; + } + index = index.wrapping_add(1); + } + self.end = index; + + self + } +} + +/// A wrapper `Fetch` type for tuple query terms that caches the last seen valid chunk. +#[doc(hidden)] +pub struct TupleFetch<'w, T: WorldQuery> { + /// the `Fetch` type of the query term + pub(crate) fetch: T::Fetch<'w>, + // the `end` value of the chunk last found in `find_table_chunk`/`find_archetype_chunk` + chunk_end: u32, +} + +impl<'w, T: WorldQuery> Clone for TupleFetch<'w, T> { + fn clone(&self) -> Self { + Self { + fetch: self.fetch.clone(), + chunk_end: self.chunk_end, + } + } } macro_rules! impl_tuple_world_query { @@ -152,32 +348,47 @@ macro_rules! impl_tuple_world_query { clippy::unused_unit, reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." )] + #[allow( + unused_mut, + reason = "Zero-length tuples will generate some function bodies that don't mutate their parameters; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." + )] $(#[$meta])* /// SAFETY: /// `fetch` accesses are the conjunction of the subqueries' accesses /// This is sound because `update_component_access` adds accesses according to the implementations of all the subqueries. /// `update_component_access` adds all `With` and `Without` filters from the subqueries. /// This is sound because `matches_component_set` always returns `false` if any the subqueries' implementations return `false`. + /// `lookahead_table` and `lookahead_archetype` always return a subset of `rows`/`indices` + /// This is sound because each `lookahead_table` and `lookahead_archetype` each unsafe impl<$($name: WorldQuery),*> WorldQuery for ($($name,)*) { - type Fetch<'w> = ($($name::Fetch<'w>,)*); + type Fetch<'w> = ($(TupleFetch<'w, $name>,)*); type State = ($($name::State,)*); fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> { let ($($name,)*) = fetch; ($( - $name::shrink_fetch($name), + TupleFetch { + fetch: $name::shrink_fetch($name.fetch), + chunk_end: $name.chunk_end + }, )*) } #[inline] unsafe fn init_fetch<'w, 's>(world: UnsafeWorldCell<'w>, state: &'s Self::State, last_run: Tick, this_run: Tick) -> Self::Fetch<'w> { let ($($name,)*) = state; - // SAFETY: The invariants are upheld by the caller. - ($(unsafe { $name::init_fetch(world, $name, last_run, this_run) },)*) + ($( + TupleFetch { + // SAFETY: The invariants are upheld by the caller. + fetch: unsafe { $name::init_fetch(world, $name, last_run, this_run) }, + chunk_end: 0, + }, + )*) } const IS_DENSE: bool = true $(&& $name::IS_DENSE)*; + const IS_ARCHETYPAL: bool = true $(&& $name::IS_ARCHETYPAL)*; #[inline] unsafe fn set_archetype<'w, 's>( @@ -188,16 +399,22 @@ macro_rules! impl_tuple_world_query { ) { let ($($name,)*) = fetch; let ($($state,)*) = state; - // SAFETY: The invariants are upheld by the caller. - $(unsafe { $name::set_archetype($name, $state, archetype, table); })* + $({ + // SAFETY: The invariants are upheld by the caller. + unsafe { $name::set_archetype(&mut $name.fetch, $state, archetype, table); }; + $name.chunk_end = 0; + })* } #[inline] unsafe fn set_table<'w, 's>(fetch: &mut Self::Fetch<'w>, state: &'s Self::State, table: &'w Table) { let ($($name,)*) = fetch; let ($($state,)*) = state; - // SAFETY: The invariants are upheld by the caller. - $(unsafe { $name::set_table($name, $state, table); })* + $({ + // SAFETY: The invariants are upheld by the caller. + unsafe { $name::set_table(&mut $name.fetch, $state, table); } + $name.chunk_end = 0; + })* } @@ -216,6 +433,109 @@ macro_rules! impl_tuple_world_query { let ($($name,)*) = state; true $(&& $name::matches_component_set($name, set_contains_id))* } + + #[inline] + unsafe fn find_table_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + table_entities: &[Entity], + mut rows: Range, + ) -> Range { + if Self::IS_ARCHETYPAL || rows.is_empty() { + rows + } else { + let mut chunk = rows.clone(); + let ($($name,)*) = fetch; + let ($($state,)*) = state; + loop { + $(if !$name::IS_ARCHETYPAL { // help the compiler out to remove unnecessary terms + let mut term_chunk = chunk.start..$name.chunk_end; + // if the term chunk is empty, we've gone past what what previously + // cached and need to refresh. + if term_chunk.is_empty() { + // SAFETY: chunk.start..rows.end is always a subset of `rows` + term_chunk = unsafe { $name::find_table_chunk($state, &mut $name.fetch, table_entities, chunk.start..rows.end) }; + $name.chunk_end = term_chunk.end; + } + // if the term chunk is empty even after refreshing the cache, there's + // no more matches in this table. + if term_chunk.is_empty() { + return term_chunk; + } + chunk.start = term_chunk.start; + chunk.end = chunk.end.min(term_chunk.end); + if chunk.is_empty() { + // the current chunk can't continue matching, so we + // have to advance its end to the next feasible chunk + chunk.end = term_chunk.end; + continue; + } + })* + + return chunk; + } + } + } + + #[inline] + unsafe fn find_archetype_chunk( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + archetype_entities: &[ArchetypeEntity], + mut indices: Range, + ) -> Range { + if Self::IS_ARCHETYPAL || indices.is_empty() { + indices + } else { + let mut chunk = indices.clone(); + let ($($name,)*) = fetch; + let ($($state,)*) = state; + loop { + $(if !$name::IS_ARCHETYPAL { // help the compiler out to remove unnecessary terms + let mut term_chunk = chunk.start..$name.chunk_end; + // if the term chunk is empty, we've gone past what what previously + // cached and need to refresh. + if term_chunk.is_empty() { + // SAFETY: chunk.start..indices.end is always a subset of `indices` + term_chunk = unsafe { $name::find_archetype_chunk($state, &mut $name.fetch, archetype_entities, chunk.start..indices.end) }; + $name.chunk_end = term_chunk.end; + } + // if the term chunk is empty even after refreshing the cache, there's + // no more matches in this table. + if term_chunk.is_empty() { + return term_chunk; + } + chunk.start = term_chunk.start; + chunk.end = chunk.end.min(term_chunk.end); + if chunk.is_empty() { + // the current chunk can't continue matching, so we + // have to advance its end to the next feasible chunk + chunk.end = term_chunk.end; + continue; + } + })* + + return chunk; + } + } + } + + #[inline] + unsafe fn matches( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + entity: Entity, + table_row: TableRow, + ) -> bool { + if Self::IS_ARCHETYPAL { + true + } else { + let ($($name,)*) = fetch; + let ($($state,)*) = state; + // SAFETY: invariants are upheld by the caller. + true $(&& unsafe { $name::matches($state, &mut $name.fetch, entity, table_row) })* + } + } } }; } diff --git a/crates/bevy_ecs/src/storage/table/mod.rs b/crates/bevy_ecs/src/storage/table/mod.rs index 846fba5ad57e3..9f3418e766836 100644 --- a/crates/bevy_ecs/src/storage/table/mod.rs +++ b/crates/bevy_ecs/src/storage/table/mod.rs @@ -97,7 +97,7 @@ impl TableId { /// [`Archetype`]: crate::archetype::Archetype /// [`Archetype::entity_table_row`]: crate::archetype::Archetype::entity_table_row /// [`Archetype::table_id`]: crate::archetype::Archetype::table_id -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] #[repr(transparent)] pub struct TableRow(NonMaxU32); diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index b8c29d448adbf..519f4bd4993bb 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -60,7 +60,7 @@ use core::{ /// ## Component access /// /// You can fetch an entity's component by specifying a reference to that component in the query's data parameter: -/// +//query.rs/ /// ``` /// # use bevy_ecs::prelude::*; /// # @@ -1581,18 +1581,26 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { D::set_archetype(&mut fetch, &self.state.fetch_state, archetype, table); F::set_archetype(&mut filter, &self.state.filter_state, archetype, table); - if F::filter_fetch( - &self.state.filter_state, - &mut filter, - entity, - location.table_row, - ) && let Some(item) = D::fetch( + let matches_fetch = D::matches( &self.state.fetch_state, &mut fetch, entity, location.table_row, - ) { - Ok(item) + ); + let matches_filter = F::matches( + &self.state.filter_state, + &mut filter, + entity, + location.table_row, + ); + + if matches_fetch && matches_filter { + Ok(D::fetch( + &self.state.fetch_state, + &mut fetch, + entity, + location.table_row, + )) } else { Err(QueryEntityError::QueryDoesNotMatch( entity, @@ -2045,7 +2053,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// This is equivalent to `self.iter().count()` but may be more efficient in some cases. /// - /// If [`D::IS_ARCHETYPAL`](QueryData::IS_ARCHETYPAL) && [`F::IS_ARCHETYPAL`](QueryFilter::IS_ARCHETYPAL) is `true`, + /// If [`D::IS_ARCHETYPAL`](crate::query::WorldQuery::IS_ARCHETYPAL) && [`F::IS_ARCHETYPAL`](crate::query::WorldQuery::IS_ARCHETYPAL) is `true`, /// this will do work proportional to the number of matched archetypes or tables, but will not iterate each entity. /// If it is `false`, it will have to do work for each entity. /// diff --git a/crates/bevy_ecs/src/world/unsafe_world_cell.rs b/crates/bevy_ecs/src/world/unsafe_world_cell.rs index cae3f5d333db7..f0d704771355b 100644 --- a/crates/bevy_ecs/src/world/unsafe_world_cell.rs +++ b/crates/bevy_ecs/src/world/unsafe_world_cell.rs @@ -1012,10 +1012,16 @@ impl<'w> UnsafeEntityCell<'w> { }; // SAFETY: Archetype and table are from the same world used to initialize state and fetch. // Table corresponds to archetype. State is the same state used to init fetch above. - unsafe { Q::set_archetype(&mut fetch, &state, archetype, table) } - // SAFETY: Called after set_archetype above. Entity and location are guaranteed to exist. - let item = unsafe { Q::fetch(&state, &mut fetch, self.id(), location.table_row) }; - item.map(Q::release_state) + unsafe { Q::set_archetype(&mut fetch, &state, archetype, table) }; + // SAFETY: + // - set_archetype was called previously + // - table_row is in range of the current archetype + unsafe { Q::matches(&state, &mut fetch, self.id(), location.table_row) } + .then(|| + // SAFETY: Called after set_archetype and above. Entity and location are + // guaranteed to exist, and this is only called if `matches` returns `true`. + unsafe { Q::fetch(&state, &mut fetch, self.id(), location.table_row) }) + .map(Q::release_state) .ok_or(QueryAccessError::EntityDoesNotMatch) } else { Err(QueryAccessError::EntityDoesNotMatch) diff --git a/crates/bevy_render/src/sync_world.rs b/crates/bevy_render/src/sync_world.rs index 6a1a1fa813a21..d43e26cd498e2 100644 --- a/crates/bevy_render/src/sync_world.rs +++ b/crates/bevy_render/src/sync_world.rs @@ -315,6 +315,7 @@ mod render_entities_world_query_impls { } const IS_DENSE: bool = <&'static RenderEntity as WorldQuery>::IS_DENSE; + const IS_ARCHETYPAL: bool = <&'static RenderEntity as WorldQuery>::IS_ARCHETYPAL; #[inline] unsafe fn set_archetype<'w, 's>( @@ -357,13 +358,23 @@ mod render_entities_world_query_impls { ) -> bool { <&RenderEntity as WorldQuery>::matches_component_set(&state, set_contains_id) } + + #[inline] + unsafe fn matches( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + entity: Entity, + table_row: TableRow, + ) -> bool { + // SAFETY: invariants are upheld by caller + unsafe { <&RenderEntity as WorldQuery>::matches(state, fetch, entity, table_row) } + } } // SAFETY: Component access of Self::ReadOnly is a subset of Self. // Self::ReadOnly matches exactly the same archetypes/tables as Self. unsafe impl QueryData for RenderEntity { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = <&MainEntity as QueryData>::IS_ARCHETYPAL; type ReadOnly = RenderEntity; type Item<'w, 's> = Entity; @@ -379,11 +390,9 @@ mod render_entities_world_query_impls { fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: defers to the `&T` implementation, with T set to `RenderEntity`. - let component = - unsafe { <&RenderEntity as QueryData>::fetch(state, fetch, entity, table_row) }; - component.map(RenderEntity::id) + unsafe { <&RenderEntity as QueryData>::fetch(state, fetch, entity, table_row).id() } } fn iter_access( @@ -430,6 +439,7 @@ mod render_entities_world_query_impls { } const IS_DENSE: bool = <&'static MainEntity as WorldQuery>::IS_DENSE; + const IS_ARCHETYPAL: bool = <&'static MainEntity as WorldQuery>::IS_ARCHETYPAL; #[inline] unsafe fn set_archetype<'w, 's>( @@ -472,13 +482,23 @@ mod render_entities_world_query_impls { ) -> bool { <&MainEntity as WorldQuery>::matches_component_set(&state, set_contains_id) } + + #[inline] + unsafe fn matches( + state: &Self::State, + fetch: &mut Self::Fetch<'_>, + entity: Entity, + table_row: TableRow, + ) -> bool { + // SAFETY: invariants are upheld by caller + unsafe { <&MainEntity as WorldQuery>::matches(state, fetch, entity, table_row) } + } } // SAFETY: Component access of Self::ReadOnly is a subset of Self. // Self::ReadOnly matches exactly the same archetypes/tables as Self. unsafe impl QueryData for MainEntity { const IS_READ_ONLY: bool = true; - const IS_ARCHETYPAL: bool = <&MainEntity as QueryData>::IS_ARCHETYPAL; type ReadOnly = MainEntity; type Item<'w, 's> = Entity; @@ -494,11 +514,9 @@ mod render_entities_world_query_impls { fetch: &mut Self::Fetch<'w>, entity: Entity, table_row: TableRow, - ) -> Option> { + ) -> Self::Item<'w, 's> { // SAFETY: defers to the `&T` implementation, with T set to `MainEntity`. - let component = - unsafe { <&MainEntity as QueryData>::fetch(state, fetch, entity, table_row) }; - component.map(MainEntity::id) + unsafe { <&MainEntity as QueryData>::fetch(state, fetch, entity, table_row).id() } } fn iter_access( diff --git a/release-content/migration-guides/query_lookahead.md b/release-content/migration-guides/query_lookahead.md new file mode 100644 index 0000000000000..7dcc8c04e7c3a --- /dev/null +++ b/release-content/migration-guides/query_lookahead.md @@ -0,0 +1,15 @@ +--- +title: "WorldQuery, QueryData, and QueryFilter trait methods rearranged" +pull_requests: [22500] +--- + +- `QueryData::IS_ARCHETYPAL`, `QueryFilter::IS_ARCHETYPAL` have been moved to `WorldQuery`. +- `QueryFilter::filter_fetch` has been moved to `WorldQuery::matches`. +- `QueryData::fetch` now returns `QueryData::Item` directly, without an `Option` + wrapper. Instead, implementations must provide `WorldQuery::matches` to + decide whether or not a row matches the query. `QueryData::fetch` can use + the return value of `matches` as a safety guarantee to avoid double-checking + a condition, for example with `Option::unwrap_unchecked`. +- `WorldQuery` has a few new methods: `find_table_chunk` and `find_archetype_chunk`. + These are unnecessary for most implementations, but add some optimization + opportunities. See module docs for more info.