diff --git a/benches/benches/bevy_ecs/index/index_iter_indexed.rs b/benches/benches/bevy_ecs/index/index_iter_indexed.rs index 22dc56f8989a4..f368d3f9839ec 100644 --- a/benches/benches/bevy_ecs/index/index_iter_indexed.rs +++ b/benches/benches/bevy_ecs/index/index_iter_indexed.rs @@ -10,8 +10,8 @@ const SPAWNS: usize = 1_000_000; struct Planet(u16); fn find_planet_zeroes_indexed(query: QueryByIndex) { - let mut query = query.at(&Planet(0)); - for planet in query.query().iter() { + let query = query.at(&Planet(0)); + for planet in query.iter() { let _ = black_box(planet); } } diff --git a/benches/benches/bevy_ecs/index/index_update_indexed.rs b/benches/benches/bevy_ecs/index/index_update_indexed.rs index 002deff143b0f..9ee08deb3387a 100644 --- a/benches/benches/bevy_ecs/index/index_update_indexed.rs +++ b/benches/benches/bevy_ecs/index/index_update_indexed.rs @@ -16,8 +16,8 @@ fn increment_planet_zeroes_indexed( let target = Planet(*local); let next_planet = Planet(target.0 + 1); - let mut query = query.at(&target); - for (entity, _planet) in query.query().iter() { + let query = query.at(&target); + for (entity, _planet) in query.iter() { commands.entity(entity).insert(next_planet); } diff --git a/crates/bevy_app/src/app.rs b/crates/bevy_app/src/app.rs index 6fb7e609ffcdb..aaf50b5884d1e 100644 --- a/crates/bevy_app/src/app.rs +++ b/crates/bevy_app/src/app.rs @@ -1357,8 +1357,7 @@ impl App { /// }); /// /// fn find_red_fans(mut query: QueryByIndex) { - /// let mut lens = query.at(&FavoriteColor::Red); - /// for entity in lens.query().iter() { + /// for entity in &query.at(&FavoriteColor::Red) { /// println!("{entity:?} likes the color Red!"); /// } /// } diff --git a/crates/bevy_ecs/src/index/mod.rs b/crates/bevy_ecs/src/index/mod.rs index 32f1c430dc0c8..d83c53525c191 100644 --- a/crates/bevy_ecs/src/index/mod.rs +++ b/crates/bevy_ecs/src/index/mod.rs @@ -92,7 +92,7 @@ //! fn get_earthlings(mut query: QueryByIndex) { //! let mut earthlings = query.at(&Planet("Earth")); //! -//! for earthling in &earthlings.query() { +//! for earthling in &earthlings { //! // ... //! } //! } diff --git a/crates/bevy_ecs/src/index/query_by_index.rs b/crates/bevy_ecs/src/index/query_by_index.rs index ac8522ee9b840..49568644b90f2 100644 --- a/crates/bevy_ecs/src/index/query_by_index.rs +++ b/crates/bevy_ecs/src/index/query_by_index.rs @@ -4,10 +4,11 @@ use crate::{ archetype::Archetype, component::{ComponentId, Immutable, Tick}, prelude::Component, - query::{QueryBuilder, QueryData, QueryFilter, QueryState, With}, - system::{Query, QueryLens, Res, SystemMeta, SystemParam}, + query::{QueryData, QueryFilter, QueryState, With}, + system::{init_query_param, Query, Res, SystemMeta, SystemParam}, world::{unsafe_world_cell::UnsafeWorldCell, World}, }; +use bevy_platform_support::sync::Arc; use super::Index; @@ -32,19 +33,17 @@ use super::Index; /// # /// # world.flush(); /// -/// fn find_all_player_one_entities(by_player: QueryByIndex) { -/// let mut lens = by_player.at(&Player(0)); -/// -/// for entity in lens.query().iter() { +/// fn find_all_player_one_entities(mut by_player: QueryByIndex) { +/// for entity in by_player.at(&Player(0)).iter() { /// println!("{entity:?} belongs to Player 1!"); /// } /// # assert_eq!(( -/// # by_player.at(&Player(0)).query().iter().count(), -/// # by_player.at(&Player(1)).query().iter().count(), -/// # by_player.at(&Player(2)).query().iter().count(), -/// # by_player.at(&Player(3)).query().iter().count(), -/// # by_player.at(&Player(4)).query().iter().count(), -/// # by_player.at(&Player(5)).query().iter().count(), +/// # by_player.at(&Player(0)).iter().count(), +/// # by_player.at(&Player(1)).iter().count(), +/// # by_player.at(&Player(2)).iter().count(), +/// # by_player.at(&Player(3)).iter().count(), +/// # by_player.at(&Player(4)).iter().count(), +/// # by_player.at(&Player(5)).iter().count(), /// # ), (1, 2, 3, 4, 5, 6)); /// } /// # world.run_system_cached(find_all_player_one_entities); @@ -57,16 +56,17 @@ pub struct QueryByIndex< F: QueryFilter + 'static = (), > { world: UnsafeWorldCell<'world>, - state: &'state QueryByIndexState, + empty_query_state: &'state QueryState)>, + query_states: &'state [QueryState)>], last_run: Tick, this_run: Tick, index: Res<'world, Index>, } -impl, D: QueryData, F: QueryFilter> - QueryByIndex<'_, '_, C, D, F> +impl<'s, C: Component, D: QueryData, F: QueryFilter> + QueryByIndex<'_, 's, C, D, F> { - /// Return a [`QueryLens`] returning entities with a component `C` of the provided value. + /// Return a [`Query`] only returning entities with a component `C` of the provided value. /// /// # Examples /// @@ -83,85 +83,71 @@ impl, D: QueryData, F: QueryFilter> /// /// world.add_index(IndexOptions::::default()); /// - /// fn find_red_fans(mut by_color: QueryByIndex) { - /// let mut lens = by_color.at(&FavoriteColor::Red); - /// - /// for entity in lens.query().iter() { + /// fn find_red_fans(mut query: QueryByIndex) { + /// for entity in query.at_mut(&FavoriteColor::Red).iter_mut() { /// println!("{entity:?} likes the color Red!"); /// } /// } /// ``` - pub fn at_mut(&mut self, value: &C) -> QueryLens<'_, D, (F, With)> - where - QueryState)>: Clone, - { - let state = self.state.primary_query_state.clone(); - + pub fn at_mut(&mut self, value: &C) -> Query<'_, 's, D, (F, With)> { // SAFETY: We have registered all of the query's world accesses, // so the caller ensures that `world` has permission to access any // world data that the query needs. unsafe { - QueryLens::new( + Query::new( self.world, - self.filter_for_value(value, state), + self.state_for_value(value), self.last_run, self.this_run, ) } } - /// Return a read-only [`QueryLens`] returning entities with a component `C` of the provided value. - pub fn at(&self, value: &C) -> QueryLens<'_, D::ReadOnly, (F, With)> - where - QueryState)>: Clone, - { - let state = self.state.primary_query_state.as_readonly().clone(); - + /// Return a read-only [`Query`] only returning entities with a component `C` of the provided value. + /// + /// # Examples + /// + /// ```rust + /// # use bevy_ecs::prelude::*; + /// # let mut world = World::new(); + /// #[derive(Component, PartialEq, Eq, Hash, Clone)] + /// #[component(immutable)] + /// enum FavoriteColor { + /// Red, + /// Green, + /// Blue, + /// } + /// + /// world.add_index::(); + /// + /// fn find_red_fans(query: QueryByIndex) { + /// for entity in query.at(&FavoriteColor::Red).iter() { + /// println!("{entity:?} likes the color Red!"); + /// } + /// } + /// ``` + pub fn at(&self, value: &C) -> Query<'_, 's, D::ReadOnly, (F, With)> { // SAFETY: We have registered all of the query's world accesses, // so the caller ensures that `world` has permission to access any // world data that the query needs. unsafe { - QueryLens::new( + Query::new( self.world, - self.filter_for_value(value, state), + self.state_for_value(value).as_readonly(), self.last_run, self.this_run, ) } } - fn filter_for_value( - &self, - value: &C, - mut state: QueryState, - ) -> QueryState { - match self.index.mapping.get(value) { - Some(index) => { - state = (0..self.index.markers.len()) - .map(|i| (i, 1 << i)) - .take_while(|&(_, mask)| mask <= self.index.slots.len()) - .map(|(i, mask)| { - if index & mask > 0 { - &self.state.with_states[i] - } else { - &self.state.without_states[i] - } - }) - .fold(state, |state, filter| { - state.join_filtered(self.world, filter) - }); - } - None => { - // Create a no-op filter by joining two conflicting filters together. - let filter = &self.state.with_states[0]; - state = state.join_filtered(self.world, filter); - - let filter = &self.state.without_states[0]; - state = state.join_filtered(self.world, filter); - } - } - - state + fn state_for_value(&self, value: &C) -> &'s QueryState)> { + let index = self.index.mapping.get(value); + // If the value was not found in the mapping, there are no matching entities and we can use the empty state. + // If the value was found but was out of range, then we have not seen any archetypes matching it yet, + // so we can still use the empty state. + index + .and_then(|index| self.query_states.get(index)) + .unwrap_or(self.empty_query_state) } } @@ -170,27 +156,27 @@ impl, D: QueryData, F: QueryFilter> /// change at any time to better suit [`QueryByIndex`]s needs. #[doc(hidden)] pub struct QueryByIndexState, D: QueryData, F: QueryFilter> { - primary_query_state: QueryState)>, + /// A `QueryState` for the underlying query that never match any archetypes. + empty_query_state: QueryState)>, + /// A list of `QueryState`s that each match the archetypes for a single index value. + query_states: Vec)>>, + /// The `SystemParam::State` for a `Res>` parameter. index_state: ComponentId, - - // TODO: Instead of storing 1 QueryState per marker component, it would be nice to instead - // track all marker components and then "somehow" create a `With`/`Without` - // query filter from that. - /// A list of [`QueryState`]s which each include a filter condition of `With`. - /// Since the marking component is dynamic, it is missing from the type signature of the state. - /// Note that this includes `With` purely for communicative purposes, `With` is a - /// strict subset of `With`. - with_states: Vec>>, - /// A list of [`QueryState`]s which each include a filter condition of `Without`. - /// Since the marking component is dynamic, it is missing from the type signature of the state. - /// Note that this includes `With` to limit the scope of this `QueryState`. - without_states: Vec>>, // No, With is not a typo + /// A copy of the markers from the `Index`. + /// We need these in `new_archetype`, but cannot access the `Index` then. + markers: Arc<[ComponentId]>, } -impl, D: QueryData + 'static, F: QueryFilter + 'static> - QueryByIndexState +// SAFETY: We rely on the known-safe implementations of `SystemParam` for `Res` and `Query`. +unsafe impl, D: QueryData + 'static, F: QueryFilter + 'static> + SystemParam for QueryByIndex<'_, '_, C, D, F> +where + QueryState)>: Clone, { - fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self { + type State = QueryByIndexState; + type Item<'w, 's> = QueryByIndex<'w, 's, C, D, F>; + + fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State { let Some(index) = world.get_resource::>() else { panic!( "Index not setup prior to usage. Please call `app.add_index(IndexOptions::<{}>::default())` during setup", @@ -198,75 +184,19 @@ impl, D: QueryData + 'static, F: QueryFilte ); }; - let ids = index.markers.clone(); + let markers = index.markers.clone(); - let primary_query_state = - )> as SystemParam>::init_state(world, system_meta); + // Start with an uninitialized `QueryState`. + // Any existing archetypes in the world will be populated when the system calls `new_archetype` during initialization. + let empty_query_state = QueryState::new_uninitialized(world); + init_query_param(world, system_meta, &empty_query_state); let index_state = > as SystemParam>::init_state(world, system_meta); - - let with_states = ids - .iter() - .map(|&id| QueryBuilder::new(world).with_id(id).build()) - .collect::>(); - - let without_states = ids - .iter() - .map(|&id| QueryBuilder::new(world).without_id(id).build()) - .collect::>(); - - Self { - primary_query_state, + QueryByIndexState { + empty_query_state, + query_states: Vec::new(), index_state, - without_states, - with_states, - } - } - - unsafe fn new_archetype(&mut self, archetype: &Archetype, system_meta: &mut SystemMeta) { - )> as SystemParam>::new_archetype( - &mut self.primary_query_state, - archetype, - system_meta, - ); - - for state in self - .with_states - .iter_mut() - .chain(self.without_states.iter_mut()) - { - > as SystemParam>::new_archetype(state, archetype, system_meta); - } - } - - #[inline] - unsafe fn validate_param(&self, system_meta: &SystemMeta, world: UnsafeWorldCell) -> bool { - let mut valid = true; - - valid &= )> as SystemParam>::validate_param( - &self.primary_query_state, - system_meta, - world, - ); - valid &= - > as SystemParam>::validate_param(&self.index_state, system_meta, world); - - for state in self.with_states.iter().chain(self.without_states.iter()) { - valid &= > as SystemParam>::validate_param(state, system_meta, world); + markers, } - - valid - } -} - -// SAFETY: We rely on the known-safe implementations of `SystemParam` for `Res` and `Query`. -unsafe impl, D: QueryData + 'static, F: QueryFilter + 'static> - SystemParam for QueryByIndex<'_, '_, C, D, F> -{ - type State = QueryByIndexState; - type Item<'w, 's> = QueryByIndex<'w, 's, C, D, F>; - - fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State { - Self::State::init_state(world, system_meta) } unsafe fn new_archetype( @@ -274,7 +204,38 @@ unsafe impl, D: QueryData + 'static, F: Que archetype: &Archetype, system_meta: &mut SystemMeta, ) { - Self::State::new_archetype(state, archetype, system_meta); + if state.empty_query_state.archetype_matches(archetype) { + // This archetype matches the overall query. + // We only want to add it to the one `QueryState` that matches the index value, + // so that it will only be returned when querying by that value. + + // We can determine the index for an archetype by looking at which + // marker components are present and adding up the relevant bits. + let indexed_markers = state.markers.iter().enumerate(); + let index = indexed_markers + .filter(|(_, &id)| archetype.contains(id)) + .map(|(i, _)| 1 << i) + .sum::(); + + // Grow the list of query states if necessary, + // cloning the empty state so that they don't match anything yet. + if index >= state.query_states.len() { + state + .query_states + .resize_with(index + 1, || state.empty_query_state.clone()); + } + + // SAFETY: + // - Caller ensures `archetype` is from the right world. + // - We called `archetype_matches` on an identical `QueryState` above. + unsafe { + state.query_states[index].new_archetype_unchecked(archetype); + state.query_states[index].update_archetype_component_access( + archetype, + system_meta.archetype_component_access_mut(), + ); + } + } } #[inline] @@ -283,7 +244,7 @@ unsafe impl, D: QueryData + 'static, F: Que system_meta: &SystemMeta, world: UnsafeWorldCell, ) -> bool { - Self::State::validate_param(state, system_meta, world) + > as SystemParam>::validate_param(&state.index_state, system_meta, world) } unsafe fn get_param<'world, 'state>( @@ -292,8 +253,6 @@ unsafe impl, D: QueryData + 'static, F: Que world: UnsafeWorldCell<'world>, change_tick: Tick, ) -> Self::Item<'world, 'state> { - state.primary_query_state.validate_world(world.id()); - let index = > as SystemParam>::get_param( &mut state.index_state, system_meta, @@ -303,7 +262,8 @@ unsafe impl, D: QueryData + 'static, F: Que QueryByIndex { world, - state, + empty_query_state: &state.empty_query_state, + query_states: &state.query_states, last_run: system_meta.last_run, this_run: change_tick, index, diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 7ae55557f1233..e2f9b44700dbd 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -233,7 +233,7 @@ impl QueryState { /// /// `new_archetype` and its variants must be called on all of the World's archetypes before the /// state can return valid query results. - fn new_uninitialized(world: &mut World) -> Self { + pub fn new_uninitialized(world: &mut World) -> Self { let fetch_state = D::init_state(world); let filter_state = F::init_state(world); Self::from_states_uninitialized(world, fetch_state, filter_state) @@ -654,6 +654,40 @@ impl QueryState { } } + /// Returns `true` if the given `archetype` matches the query. Otherwise, returns `false`. + pub fn archetype_matches(&self, archetype: &Archetype) -> bool { + D::matches_component_set(&self.fetch_state, &|id| archetype.contains(id)) + && F::matches_component_set(&self.filter_state, &|id| archetype.contains(id)) + && self.matches_component_set(&|id| archetype.contains(id)) + } + + /// Process the given [`Archetype`] to update internal metadata about the [`Table`](crate::storage::Table)s + /// and [`Archetype`]s that are matched by this query. + /// + /// # Safety + /// `archetype` must be from the `World` this state was initialized from, + /// and must match the current query. + pub unsafe fn new_archetype_unchecked(&mut self, archetype: &Archetype) { + let archetype_index = archetype.id().index(); + if !self.matched_archetypes.contains(archetype_index) { + self.matched_archetypes.grow_and_insert(archetype_index); + if !self.is_dense { + self.matched_storage_ids.push(StorageId { + archetype_id: archetype.id(), + }); + } + } + let table_index = archetype.table_id().as_usize(); + if !self.matched_tables.contains(table_index) { + self.matched_tables.grow_and_insert(table_index); + if self.is_dense { + self.matched_storage_ids.push(StorageId { + table_id: archetype.table_id(), + }); + } + } + } + /// Process the given [`Archetype`] to update internal metadata about the [`Table`](crate::storage::Table)s /// and [`Archetype`]s that are matched by this query. /// @@ -663,28 +697,8 @@ impl QueryState { /// # Safety /// `archetype` must be from the `World` this state was initialized from. unsafe fn new_archetype_internal(&mut self, archetype: &Archetype) -> bool { - if D::matches_component_set(&self.fetch_state, &|id| archetype.contains(id)) - && F::matches_component_set(&self.filter_state, &|id| archetype.contains(id)) - && self.matches_component_set(&|id| archetype.contains(id)) - { - let archetype_index = archetype.id().index(); - if !self.matched_archetypes.contains(archetype_index) { - self.matched_archetypes.grow_and_insert(archetype_index); - if !self.is_dense { - self.matched_storage_ids.push(StorageId { - archetype_id: archetype.id(), - }); - } - } - let table_index = archetype.table_id().as_usize(); - if !self.matched_tables.contains(table_index) { - self.matched_tables.grow_and_insert(table_index); - if self.is_dense { - self.matched_storage_ids.push(StorageId { - table_id: archetype.table_id(), - }); - } - } + if self.archetype_matches(archetype) { + self.new_archetype_unchecked(archetype); true } else { false diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 362652c69d9c6..ad1413d3b149c 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -2102,33 +2102,6 @@ impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> { // uphold the safety invariants of `Query::new` unsafe { Query::new(self.world, &self.state, self.last_run, self.this_run) } } - - /// Creates a new [`QueryLens`]. - /// - /// # Panics - /// - /// This will panic if the world used to create `state` is not `world`. - /// - /// # Safety - /// - /// `QueryLens` can be used to construct a `Query` by internally calling `Query::new`. - /// All safety invariants of `Query::new` must be upheld when calling `QueryLens::new`. - #[inline] - pub(crate) unsafe fn new( - world: UnsafeWorldCell<'w>, - state: QueryState, - last_run: Tick, - this_run: Tick, - ) -> Self { - state.validate_world(world.id()); - - Self { - world, - state, - last_run, - this_run, - } - } } impl<'w, 's, Q: QueryData, F: QueryFilter> From<&'s mut QueryLens<'w, Q, F>>