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<Item: Borrow<Entity>>`).

`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<Entity>,
}

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 <deontologician@gmail.com>

Co-authored-by: devil-ira <justthecooldude@gmail.com>
This commit is contained in:
ira 2022-06-06 16:09:16 +00:00
parent b47291264b
commit 92ddfe8ad4
7 changed files with 449 additions and 8 deletions

View File

@ -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>,
{
entity_iter: I,
entities: &'w Entities,
tables: &'w Tables,
archetypes: &'w Archetypes,
fetch: QF,
filter: QueryFetch<'w, F>,
query_state: &'s QueryState<Q, F>,
}
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<Entity>,
{
/// # 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<EntityList: IntoIterator<IntoIter = I>>(
world: &'w World,
query_state: &'s QueryState<Q, F>,
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::<F>::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<Entity>,
{
type Item = QF::Item;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
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<usize>) {
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,

View File

@ -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<Entity, With<A>>, 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<Entity, With<A>>, 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);
}
}
}

View File

@ -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<Q: WorldQuery, F: WorldQuery = ()> {
@ -556,6 +556,32 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<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.
///
#[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<Entity>,
{
// 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<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
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<Entity>,
{
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<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
);
}
/// Runs `func` on each query result where the entities match.
#[inline]
pub fn many_for_each_mut<EntityList: IntoIterator>(
&mut self,
world: &mut World,
entities: EntityList,
func: impl FnMut(QueryItem<'_, Q>),
) where
EntityList::Item: Borrow<Entity>,
{
// 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<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
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 = <QueryFetch<F> as Fetch>::init(
world,
@ -866,7 +944,7 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
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<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
});
}
/// 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<EntityList: IntoIterator>(
&self,
world: &World,
entity_list: EntityList,
mut func: impl FnMut(QueryItem<'_, Q>),
last_change_tick: u32,
change_tick: u32,
) where
EntityList::Item: Borrow<Entity>,
{
// 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 =
<QueryFetch<Q> as Fetch>::init(world, &self.fetch_state, last_change_tick, change_tick);
let mut filter = <QueryFetch<F> 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.
///

View File

@ -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<Entity>,
/// }
///
/// 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<EntityList: IntoIterator>(
&self,
entities: EntityList,
) -> QueryManyIter<'_, '_, Q, ROQueryFetch<'_, Q>, F, EntityList::IntoIter>
where
EntityList::Item: Borrow<Entity>,
{
// 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<EntityList: IntoIterator>(
&self,
entities: EntityList,
) -> QueryManyIter<'_, '_, Q, QueryFetch<'_, Q>, F, EntityList::IntoIter>
where
EntityList::Item: Borrow<Entity>,
{
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<Entity>,
/// }
///
/// 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<EntityList: IntoIterator>(
&mut self,
entities: EntityList,
f: impl FnMut(QueryItem<'_, Q>),
) where
EntityList::Item: Borrow<Entity>,
{
// 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

View File

@ -0,0 +1,14 @@
use bevy_ecs::prelude::*;
#[derive(Component)]
struct A(usize);
fn system(mut query: Query<&mut A>, e: Res<Entity>) {
let mut results = Vec::new();
query.many_for_each_mut(vec![*e, *e], |a| {
// this should fail to compile
results.push(a);
});
}
fn main() {}

View File

@ -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

View File

@ -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 = <Self::IntoIter as Iterator>::Item;
type IntoIter = slice::Iter<'a, Entity>;
fn into_iter(self) -> Self::IntoIter {
self.0.iter()
}
}