diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 6f12ea557d20d..84843283cc6fb 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1,10 +1,11 @@ use crate::{ archetype::{ArchetypeId, Archetypes}, + entity::{Entities, Entity}, + prelude::World, query::{Fetch, QueryState, WorldQuery}, storage::{TableId, Tables}, - world::World, }; -use std::{marker::PhantomData, mem::MaybeUninit}; +use std::{borrow::Borrow, marker::PhantomData, mem::MaybeUninit}; use super::{QueryFetch, QueryItem, ReadOnlyFetch}; @@ -71,6 +72,113 @@ where } } +/// An [`Iterator`] over query results of a [`Query`](crate::system::Query). +/// +/// This struct is created by the [`Query::iter_many`](crate::system::Query::iter_many) method. +pub struct QueryManyIter< + 'w, + 's, + Q: WorldQuery, + QF: Fetch<'w, State = Q::State>, + F: WorldQuery, + I: Iterator, +> where + I::Item: Borrow, +{ + entity_iter: I, + entities: &'w Entities, + tables: &'w Tables, + archetypes: &'w Archetypes, + fetch: QF, + filter: QueryFetch<'w, F>, + query_state: &'s QueryState, +} + +impl<'w, 's, Q: WorldQuery, QF: Fetch<'w, State = Q::State>, F: WorldQuery, I: Iterator> + QueryManyIter<'w, 's, Q, QF, F, I> +where + I::Item: Borrow, +{ + /// # Safety + /// This does not check for mutable query correctness. To be safe, make sure mutable queries + /// have unique access to the components they query. + /// This does not validate that `world.id()` matches `query_state.world_id`. Calling this on a `world` + /// with a mismatched [`WorldId`](crate::world::WorldId) is unsound. + pub(crate) unsafe fn new>( + world: &'w World, + query_state: &'s QueryState, + entity_list: EntityList, + last_change_tick: u32, + change_tick: u32, + ) -> QueryManyIter<'w, 's, Q, QF, F, I> { + let fetch = QF::init( + world, + &query_state.fetch_state, + last_change_tick, + change_tick, + ); + let filter = QueryFetch::::init( + world, + &query_state.filter_state, + last_change_tick, + change_tick, + ); + QueryManyIter { + query_state, + entities: &world.entities, + archetypes: &world.archetypes, + tables: &world.storages.tables, + fetch, + filter, + entity_iter: entity_list.into_iter(), + } + } +} + +impl<'w, 's, Q: WorldQuery, QF: Fetch<'w, State = Q::State>, F: WorldQuery, I: Iterator> Iterator + for QueryManyIter<'w, 'w, Q, QF, F, I> +where + I::Item: Borrow, +{ + type Item = QF::Item; + + #[inline(always)] + fn next(&mut self) -> Option { + unsafe { + for entity in self.entity_iter.by_ref() { + let location = match self.entities.get(*entity.borrow()) { + Some(location) => location, + None => continue, + }; + + if !self + .query_state + .matched_archetypes + .contains(location.archetype_id.index()) + { + continue; + } + + let archetype = &self.archetypes[location.archetype_id]; + + self.fetch + .set_archetype(&self.query_state.fetch_state, archetype, self.tables); + self.filter + .set_archetype(&self.query_state.filter_state, archetype, self.tables); + if self.filter.archetype_filter_fetch(location.index) { + return Some(self.fetch.archetype_fetch(location.index)); + } + } + None + } + } + + fn size_hint(&self) -> (usize, Option) { + let (_, max_size) = self.entity_iter.size_hint(); + (0, max_size) + } +} + pub struct QueryCombinationIter<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> { tables: &'w Tables, archetypes: &'w Archetypes, diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 4e01c9e20d4fd..857fe033ab1d4 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -21,6 +21,7 @@ unsafe fn debug_checked_unreachable() -> ! { mod tests { use super::WorldQuery; use crate::prelude::{AnyOf, Entity, Or, With, Without}; + use crate::system::{IntoSystem, Query, System}; use crate::{self as bevy_ecs, component::Component, world::World}; use std::collections::HashSet; @@ -516,4 +517,44 @@ mod tests { assert_eq!(custom_param_entities, normal_entities); } } + + #[test] + fn many_entities() { + let mut world = World::new(); + world.spawn().insert_bundle((A(0), B(0))); + world.spawn().insert_bundle((A(0), B(0))); + world.spawn().insert(A(0)); + world.spawn().insert(B(0)); + { + fn system(has_a: Query>, has_a_and_b: Query<(&A, &B)>) { + assert_eq!(has_a_and_b.iter_many(&has_a).count(), 2); + } + let mut system = IntoSystem::into_system(system); + system.initialize(&mut world); + system.run((), &mut world); + } + { + fn system(has_a: Query>, mut b_query: Query<&mut B>) { + b_query.many_for_each_mut(&has_a, |mut b| { + b.0 = 1; + }); + } + let mut system = IntoSystem::into_system(system); + system.initialize(&mut world); + system.run((), &mut world); + } + { + fn system(query: Query<(Option<&A>, &B)>) { + for (maybe_a, b) in &query { + match maybe_a { + Some(_) => assert_eq!(b.0, 1), + None => assert_eq!(b.0, 0), + } + } + } + let mut system = IntoSystem::into_system(system); + system.initialize(&mut world); + system.run((), &mut world); + } + } } diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index b88407f69212f..64a80313c69f1 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -14,9 +14,9 @@ use bevy_tasks::{ComputeTaskPool, TaskPool}; #[cfg(feature = "trace")] use bevy_utils::tracing::Instrument; use fixedbitset::FixedBitSet; -use std::{fmt, ops::Deref}; +use std::{borrow::Borrow, fmt, ops::Deref}; -use super::{QueryFetch, QueryItem, ROQueryFetch, ROQueryItem}; +use super::{QueryFetch, QueryItem, QueryManyIter, ROQueryFetch, ROQueryItem}; /// Provides scoped access to a [`World`] state according to a given [`WorldQuery`] and query filter. pub struct QueryState { @@ -556,6 +556,32 @@ impl QueryState { } } + /// Returns an [`Iterator`] over the query results of a list of [`Entity`]'s. + /// + /// This can only return immutable data (mutable data will be cast to an immutable form). + /// See [`Self::many_for_each_mut`] for queries that contain at least one mutable component. + /// + #[inline] + pub fn iter_many<'w, 's, EntityList: IntoIterator>( + &'s mut self, + world: &'w World, + entities: EntityList, + ) -> QueryManyIter<'w, 's, Q, ROQueryFetch<'w, Q>, F, EntityList::IntoIter> + where + EntityList::Item: Borrow, + { + // SAFETY: query is read only + unsafe { + self.update_archetypes(world); + self.iter_many_unchecked_manual( + entities, + world, + world.last_change_tick(), + world.read_change_tick(), + ) + } + } + /// Returns an [`Iterator`] over the query results for the given [`World`]. /// /// # Safety @@ -611,6 +637,35 @@ impl QueryState { QueryIter::new(world, self, last_change_tick, change_tick) } + /// Returns an [`Iterator`] for the given [`World`] and list of [`Entity`]'s, where the last change and + /// the current change tick are given. + /// + /// # Safety + /// + /// This does not check for mutable query correctness. To be safe, make sure mutable queries + /// have unique access to the components they query. + /// this does not check for entity uniqueness + /// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world` + /// with a mismatched [`WorldId`] is unsound. + #[inline] + pub(crate) unsafe fn iter_many_unchecked_manual< + 'w, + 's, + QF: Fetch<'w, State = Q::State>, + EntityList: IntoIterator, + >( + &'s self, + entities: EntityList, + world: &'w World, + last_change_tick: u32, + change_tick: u32, + ) -> QueryManyIter<'w, 's, Q, QF, F, EntityList::IntoIter> + where + EntityList::Item: Borrow, + { + QueryManyIter::new(world, self, entities, last_change_tick, change_tick) + } + /// Returns an [`Iterator`] over all possible combinations of `K` query results for the /// given [`World`] without repetition. /// This can only be called for read-only queries. @@ -775,6 +830,29 @@ impl QueryState { ); } + /// Runs `func` on each query result where the entities match. + #[inline] + pub fn many_for_each_mut( + &mut self, + world: &mut World, + entities: EntityList, + func: impl FnMut(QueryItem<'_, Q>), + ) where + EntityList::Item: Borrow, + { + // SAFETY: query has unique world access + unsafe { + self.update_archetypes(world); + self.many_for_each_unchecked_manual( + world, + entities, + func, + world.last_change_tick(), + world.read_change_tick(), + ); + }; + } + /// Runs `func` on each query result for the given [`World`], where the last change and /// the current change tick are given. This is faster than the equivalent /// iter() method, but cannot be chained like a normal [`Iterator`]. @@ -797,7 +875,7 @@ impl QueryState { change_tick: u32, ) { // NOTE: If you are changing query iteration code, remember to update the following places, where relevant: - // QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual + // QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual let mut fetch = QF::init(world, &self.fetch_state, last_change_tick, change_tick); let mut filter = as Fetch>::init( world, @@ -866,7 +944,7 @@ impl QueryState { change_tick: u32, ) { // NOTE: If you are changing query iteration code, remember to update the following places, where relevant: - // QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual + // QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual self.task_pool .as_ref() .expect("Cannot iterate query in parallel. No ComputeTaskPool initialized.") @@ -969,6 +1047,62 @@ impl QueryState { }); } + /// Runs `func` on each query result for the given [`World`] and list of [`Entity`]'s, where the last change and + /// the current change tick are given. This is faster than the equivalent + /// iter() method, but cannot be chained like a normal [`Iterator`]. + /// + /// # Safety + /// + /// This does not check for mutable query correctness. To be safe, make sure mutable queries + /// have unique access to the components they query. + /// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world` + /// with a mismatched [`WorldId`] is unsound. + pub(crate) unsafe fn many_for_each_unchecked_manual( + &self, + world: &World, + entity_list: EntityList, + mut func: impl FnMut(QueryItem<'_, Q>), + last_change_tick: u32, + change_tick: u32, + ) where + EntityList::Item: Borrow, + { + // NOTE: If you are changing query iteration code, remember to update the following places, where relevant: + // QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual + let mut fetch = + as Fetch>::init(world, &self.fetch_state, last_change_tick, change_tick); + let mut filter = as Fetch>::init( + world, + &self.filter_state, + last_change_tick, + change_tick, + ); + + let tables = &world.storages.tables; + + for entity in entity_list.into_iter() { + let location = match world.entities.get(*entity.borrow()) { + Some(location) => location, + None => continue, + }; + + if !self + .matched_archetypes + .contains(location.archetype_id.index()) + { + continue; + } + + let archetype = &world.archetypes[location.archetype_id]; + + fetch.set_archetype(&self.fetch_state, archetype, tables); + filter.set_archetype(&self.filter_state, archetype, tables); + if filter.archetype_filter_fetch(location.index) { + func(fetch.archetype_fetch(location.index)); + } + } + } + /// Returns a single immutable query result when there is exactly one entity matching /// the query. /// diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 4fbab3b47fe88..484d119e4111d 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -3,11 +3,12 @@ use crate::{ entity::Entity, query::{ NopFetch, QueryCombinationIter, QueryEntityError, QueryFetch, QueryItem, QueryIter, - QuerySingleError, QueryState, ROQueryFetch, ROQueryItem, ReadOnlyFetch, WorldQuery, + QueryManyIter, QuerySingleError, QueryState, ROQueryFetch, ROQueryItem, ReadOnlyFetch, + WorldQuery, }, world::{Mut, World}, }; -use std::{any::TypeId, fmt::Debug}; +use std::{any::TypeId, borrow::Borrow, fmt::Debug}; /// Provides scoped access to components in a [`World`]. /// @@ -387,6 +388,56 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> { } } + /// Returns an [`Iterator`] over the query results of a list of [`Entity`]'s. + /// + /// This can only return immutable data (mutable data will be cast to an immutable form). + /// See [`Self::many_for_each_mut`] for queries that contain at least one mutable component. + /// + /// # Examples + /// ``` + /// # use bevy_ecs::prelude::*; + /// #[derive(Component)] + /// struct Counter { + /// value: i32 + /// } + /// + /// #[derive(Component)] + /// struct Friends { + /// list: Vec, + /// } + /// + /// fn system( + /// friends_query: Query<&Friends>, + /// counter_query: Query<&Counter>, + /// ) { + /// for friends in &friends_query { + /// for counter in counter_query.iter_many(&friends.list) { + /// println!("Friend's counter: {:?}", counter.value); + /// } + /// } + /// } + /// # bevy_ecs::system::assert_is_system(system); + /// ``` + #[inline] + pub fn iter_many( + &self, + entities: EntityList, + ) -> QueryManyIter<'_, '_, Q, ROQueryFetch<'_, Q>, F, EntityList::IntoIter> + where + EntityList::Item: Borrow, + { + // SAFETY: system runs without conflicts with other systems. + // same-system queries have runtime borrow checks when they conflict + unsafe { + self.state.iter_many_unchecked_manual( + entities, + self.world, + self.last_change_tick, + self.change_tick, + ) + } + } + /// Returns an [`Iterator`] over the query results. /// /// # Safety @@ -420,6 +471,29 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> { ) } + /// Returns an [`Iterator`] over the query results of a list of [`Entity`]'s. + /// + /// If you want safe mutable access to query results of a list of [`Entity`]'s. See [`Self::many_for_each_mut`]. + /// + /// # Safety + /// This allows aliased mutability and does not check for entity uniqueness. + /// You must make sure this call does not result in multiple mutable references to the same component. + /// Particular care must be taken when collecting the data (rather than iterating over it one item at a time) such as via `[Iterator::collect()]`. + pub unsafe fn iter_many_unsafe( + &self, + entities: EntityList, + ) -> QueryManyIter<'_, '_, Q, QueryFetch<'_, Q>, F, EntityList::IntoIter> + where + EntityList::Item: Borrow, + { + self.state.iter_many_unchecked_manual( + entities, + self.world, + self.last_change_tick, + self.change_tick, + ) + } + /// Runs `f` on each query result. This is faster than the equivalent iter() method, but cannot /// be chained like a normal [`Iterator`]. /// @@ -565,6 +639,55 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> { }; } + /// Calls a closure on each result of [`Query`] where the entities match. + /// # Examples + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// #[derive(Component)] + /// struct Counter { + /// value: i32 + /// } + /// + /// #[derive(Component)] + /// struct Friends { + /// list: Vec, + /// } + /// + /// fn system( + /// friends_query: Query<&Friends>, + /// mut counter_query: Query<&mut Counter>, + /// ) { + /// for friends in &friends_query { + /// counter_query.many_for_each_mut(&friends.list, |mut counter| { + /// println!("Friend's counter: {:?}", counter.value); + /// counter.value += 1; + /// }); + /// } + /// } + /// # bevy_ecs::system::assert_is_system(system); + /// ``` + #[inline] + pub fn many_for_each_mut( + &mut self, + entities: EntityList, + f: impl FnMut(QueryItem<'_, Q>), + ) where + EntityList::Item: Borrow, + { + // SAFE: system runs without conflicts with other systems. + // same-system queries have runtime borrow checks when they conflict + unsafe { + self.state.many_for_each_unchecked_manual( + self.world, + entities, + f, + self.last_change_tick, + self.change_tick, + ); + }; + } + /// Returns the query result for the given [`Entity`]. /// /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.rs new file mode 100644 index 0000000000000..e35ba6dbc57d2 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.rs @@ -0,0 +1,14 @@ +use bevy_ecs::prelude::*; + +#[derive(Component)] +struct A(usize); + +fn system(mut query: Query<&mut A>, e: Res) { + let mut results = Vec::new(); + query.many_for_each_mut(vec![*e, *e], |a| { + // this should fail to compile + results.push(a); + }); +} + +fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.stderr new file mode 100644 index 0000000000000..b37e88e932473 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.stderr @@ -0,0 +1,10 @@ +error[E0521]: borrowed data escapes outside of closure + --> tests/ui/system_query_many_for_each_mut_lifetime_safety.rs:10:9 + | +7 | let mut results = Vec::new(); + | ----------- `results` declared here, outside of the closure body +8 | query.many_for_each_mut(vec![*e, *e], |a| { + | - `a` is a reference that is only valid in the closure body +9 | // this should fail to compile +10 | results.push(a); + | ^^^^^^^^^^^^^^^ `a` escapes the closure body here diff --git a/crates/bevy_hierarchy/src/components/children.rs b/crates/bevy_hierarchy/src/components/children.rs index 96097a1da8a59..d6c1e42604504 100644 --- a/crates/bevy_hierarchy/src/components/children.rs +++ b/crates/bevy_hierarchy/src/components/children.rs @@ -4,6 +4,7 @@ use bevy_ecs::{ reflect::{ReflectComponent, ReflectMapEntities}, }; use bevy_reflect::Reflect; +use core::slice; use smallvec::SmallVec; use std::ops::Deref; @@ -41,3 +42,13 @@ impl Deref for Children { &self.0[..] } } + +impl<'a> IntoIterator for &'a Children { + type Item = ::Item; + + type IntoIter = slice::Iter<'a, Entity>; + + fn into_iter(self) -> Self::IntoIter { + self.0.iter() + } +}