[Merged by Bors] - Add methods for querying lists of entities.

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,114 @@ where
     }
 }
 
+/// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
+///
+/// This struct is created by the [`Query::many_iter`](crate::system::Query::many_iter) and
+/// [`Query::many_iter_mut`](crate::system::Query::many_iter_mut) methods.
+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;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        unsafe {
+            for entity in self.entity_iter.by_ref() {
+                let location = if let Some(location) = self.entities.get(*entity.borrow()) {
+                    location
+                } else {
+                    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,

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_iter() {
+        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.many_iter(&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>) {
+                for mut b in b_query.many_iter_mut(&has_a) {
+                    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);
+        }
+    }
 }

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<Q: WorldQuery, F: WorldQuery = ()> {
@@ -611,6 +611,34 @@ 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 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.
@@ -797,7 +825,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 +894,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.")

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`].
 ///
@@ -1127,6 +1128,99 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
                 .is_ok()
         }
     }
+
+    /// 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_iter_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.many_iter(&friends.list) {
+    ///             println!("Friend's counter: {:?}", counter.value);
+    ///         }
+    ///     }
+    /// }
+    /// # bevy_ecs::system::assert_is_system(system);
+    /// ```
+    pub fn many_iter<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 of a list of [`Entity`]'s.
+    ///
+    /// # 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 {
+    ///         for mut counter in counter_query.many_iter_mut(&friends.list) {
+    ///             println!("Friend's counter: {:?}", counter.value);
+    ///             counter.value += 1;
+    ///         }
+    ///     }
+    /// }
+    /// ```
+    pub fn many_iter_mut<EntityList: IntoIterator>(
+        &mut self,
+        entities: EntityList,
+    ) -> QueryManyIter<'_, '_, Q, QueryFetch<'_, 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,
+            )
+        }
+    }
 }
 
 impl<'w, 's, Q: WorldQuery, F: WorldQuery> IntoIterator for &'w Query<'_, 's, Q, F> {

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