From d924c05c6a11609d0d1487627e41ece33bbb7559 Mon Sep 17 00:00:00 2001 From: ira Date: Mon, 6 Jun 2022 16:09:16 +0000 Subject: [PATCH] Add methods for querying lists of entities. (#4879) # Objective Improve querying ergonomics around collections and iterators of entities. Example how queries over Children might be done currently. ```rust fn system(foo_query: Query<(&Foo, &Children)>, bar_query: Query<(&Bar, &Children)>) { for (foo, children) in &foo_query { for child in children.iter() { if let Ok((bar, children)) = bar_query.get(*child) { for child in children.iter() { if let Ok((foo, children)) = foo_query.get(*child) { // D: } } } } } } ``` Answers #4868 Partially addresses #4864 Fixes #1470 ## Solution Based on the great work by @deontologician in #2563 Added `iter_many` and `many_for_each_mut` to `Query`. These take a list of entities (Anything that implements `IntoIterator>`). `iter_many` returns a `QueryManyIter` iterator over immutable results of a query (mutable data will be cast to an immutable form). `many_for_each_mut` calls a closure for every result of the query, ensuring not aliased mutability. This iterator goes over the list of entities in order and returns the result from the query for it. Skipping over any entities that don't match the query. Also added `unsafe fn iter_many_unsafe`. ### Examples ```rust #[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 { for counter in counter_query.iter_many(&friends.list) { println!("Friend's counter: {:?}", counter.value); } counter_query.many_for_each_mut(&friends.list, |mut counter| { counter.value += 1; println!("Friend's counter: {:?}", counter.value); }); } } ``` Here's how example in the Objective section can be written with this PR. ```rust fn system(foo_query: Query<(&Foo, &Children)>, bar_query: Query<(&Bar, &Children)>) { for (foo, children) in &foo_query { for (bar, children) in bar_query.iter_many(children) { for (foo, children) in foo_query.iter_many(children) { // :D } } } } ``` ## Additional changes Implemented `IntoIterator` for `&Children` because why not. ## Todo - Bikeshed! Co-authored-by: deontologician Co-authored-by: devil-ira --- crates/bevy_ecs/src/query/iter.rs | 112 +++++++++++++- crates/bevy_ecs/src/query/mod.rs | 41 +++++ crates/bevy_ecs/src/query/state.rs | 142 +++++++++++++++++- crates/bevy_ecs/src/system/query.rs | 127 +++++++++++++++- ...query_many_for_each_mut_lifetime_safety.rs | 14 ++ ...y_many_for_each_mut_lifetime_safety.stderr | 10 ++ .../bevy_hierarchy/src/components/children.rs | 11 ++ 7 files changed, 449 insertions(+), 8 deletions(-) create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.rs create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.stderr diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 6f12ea557d20d2..84843283cc6fb9 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 4e01c9e20d4fdc..857fe033ab1d47 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 b88407f69212fd..64a80313c69f1a 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 4fbab3b47fe88c..484d119e4111db 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 00000000000000..e35ba6dbc57d2b --- /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 00000000000000..b37e88e9324732 --- /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 96097a1da8a59b..d6c1e426045043 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() + } +}