[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,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,

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<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);
+        }
+    }
 }

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 = ()> {
@@ -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.
     ///