Skip to content

Commit

Permalink
Add methods for querying lists of entities. (#4879)
Browse files Browse the repository at this point in the history
# 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>
  • Loading branch information
irate-devil and irate-devil committed Jun 6, 2022
1 parent 292cd03 commit 94d1d23
Show file tree
Hide file tree
Showing 7 changed files with 449 additions and 8 deletions.
112 changes: 110 additions & 2 deletions crates/bevy_ecs/src/query/iter.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand Down Expand Up @@ -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,
Expand Down
41 changes: 41 additions & 0 deletions crates/bevy_ecs/src/query/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
}
}
}
142 changes: 138 additions & 4 deletions crates/bevy_ecs/src/query/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ()> {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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`].
Expand All @@ -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,
Expand Down Expand Up @@ -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.")
Expand Down Expand Up @@ -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.
///
Expand Down

0 comments on commit 94d1d23

Please sign in to comment.