Add iter_join_map. Refactor compile fail test.

crates/bevy_ecs/src/query/iter.rs

@@ -208,6 +208,140 @@ where
 {
 }
 
+/// An [`Iterator`] over [`Query`](crate::system::Query) results a list mapped to [`Entity`]'s and the list items.
+///
+/// This struct is created by the [`Query::iter_join_map`](crate::system::Query::iter_join_map) and [`Query::iter_join_map_mut`](crate::system::Query::iter_join_map_mut) methods.
+pub struct QueryJoinMapIter<'w, 's, Q: WorldQuery, F: WorldQuery, I: Iterator, MapFn>
+where
+    MapFn: FnMut(&I::Item) -> Entity,
+{
+    list: I,
+    map_f: MapFn,
+    entities: &'w Entities,
+    tables: &'w Tables,
+    archetypes: &'w Archetypes,
+    fetch: QueryFetch<'w, Q>,
+    filter: QueryFetch<'w, F>,
+    query_state: &'s QueryState<Q, F>,
+}
+
+impl<'w, 's, Q: WorldQuery, F: WorldQuery, I: Iterator, MapFn>
+    QueryJoinMapIter<'w, 's, Q, F, I, MapFn>
+where
+    MapFn: FnMut(&I::Item) -> 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<II: IntoIterator<IntoIter = I>>(
+        world: &'w World,
+        query_state: &'s QueryState<Q, F>,
+        last_change_tick: u32,
+        change_tick: u32,
+        entity_map: II,
+        map_f: MapFn,
+    ) -> QueryJoinMapIter<'w, 's, Q, F, I, MapFn> {
+        let fetch = Q::Fetch::init(
+            world,
+            &query_state.fetch_state,
+            last_change_tick,
+            change_tick,
+        );
+        let filter = F::Fetch::init(
+            world,
+            &query_state.filter_state,
+            last_change_tick,
+            change_tick,
+        );
+        QueryJoinMapIter {
+            query_state,
+            entities: &world.entities,
+            archetypes: &world.archetypes,
+            tables: &world.storages.tables,
+            fetch,
+            filter,
+            list: entity_map.into_iter(),
+            map_f,
+        }
+    }
+
+    /// SAFETY:
+    /// The lifetime here is not restrictive enough for Fetch with &mut access,
+    /// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple
+    /// references to the same component, leading to unique reference aliasing.
+    ///
+    /// It is always safe for immutable borrows.
+    #[inline(always)]
+    unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<(QueryItem<'w, Q>, I::Item)> {
+        for item in self.list.by_ref() {
+            let location = match self.entities.get((self.map_f)(&item)) {
+                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), item));
+            }
+        }
+        None
+    }
+
+    /// Get the next item from the iterator.
+    #[inline(always)]
+    pub fn fetch_next(&mut self) -> Option<(QueryItem<'_, Q>, I::Item)> {
+        // safety: we are limiting the returned reference to self,
+        // making sure this method cannot be called multiple times without getting rid
+        // of any previously returned unique references first, thus preventing aliasing.
+        unsafe {
+            self.fetch_next_aliased_unchecked()
+                .map(|(q_item, item)| (Q::shrink(q_item), item))
+        }
+    }
+}
+
+impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator, MapFn> Iterator
+    for QueryJoinMapIter<'w, 's, Q, F, I, MapFn>
+where
+    MapFn: FnMut(&I::Item) -> Entity,
+{
+    type Item = (QueryItem<'w, Q>, I::Item);
+
+    #[inline(always)]
+    fn next(&mut self) -> Option<Self::Item> {
+        // SAFETY: it is safe to alias for ReadOnlyWorldQuery
+        unsafe { self.fetch_next_aliased_unchecked() }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (_, max_size) = self.list.size_hint();
+        (0, max_size)
+    }
+}
+
+// This is correct as QueryJoinMapIter always returns `None` once exhausted.
+impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator, MapFn> FusedIterator
+    for QueryJoinMapIter<'w, 's, Q, F, I, MapFn>
+where
+    MapFn: FnMut(&I::Item) -> Entity,
+{
+}
+
 pub struct QueryCombinationIter<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> {
     tables: &'w Tables,
     archetypes: &'w Archetypes,

crates/bevy_ecs/src/query/state.rs

@@ -15,7 +15,7 @@ use bevy_utils::tracing::Instrument;
 use fixedbitset::FixedBitSet;
 use std::{borrow::Borrow, fmt};
 
-use super::{NopWorldQuery, QueryFetch, QueryItem, QueryManyIter, ROQueryItem};
+use super::{NopWorldQuery, QueryFetch, QueryItem, QueryJoinMapIter, QueryManyIter, ROQueryItem};
 
 /// Provides scoped access to a [`World`] state according to a given [`WorldQuery`] and query filter.
 #[repr(C)]
@@ -647,6 +647,50 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
         }
     }
 
+    /// Returns an [`Iterator`] over the inner join of the results of a query and list of items mapped to [`Entity`]'s.
+    ///
+    /// This can only be called for read-only queries, see [`Self::iter_list_mut`] for write-queries.
+    #[inline]
+    pub fn iter_join_map<'w, 's, I: IntoIterator, MapFn: FnMut(&I::Item) -> Entity>(
+        &'s mut self,
+        world: &'w World,
+        list: I,
+        map_f: MapFn,
+    ) -> QueryJoinMapIter<'w, 's, Q::ReadOnly, F::ReadOnly, I::IntoIter, MapFn> {
+        self.update_archetypes(world);
+        // SAFETY: Query has unique world access.
+        unsafe {
+            self.as_readonly().iter_join_map_unchecked_manual(
+                world,
+                world.last_change_tick(),
+                world.read_change_tick(),
+                list,
+                map_f,
+            )
+        }
+    }
+
+    /// Returns an iterator over the inner join of the results of a query and list of items mapped to [`Entity`]'s.
+    #[inline]
+    pub fn iter_join_map_mut<'w, 's, I: IntoIterator, MapFn: FnMut(&I::Item) -> Entity>(
+        &'s mut self,
+        world: &'w mut World,
+        list: I,
+        map_f: MapFn,
+    ) -> QueryJoinMapIter<'w, 's, Q, F, I::IntoIter, MapFn> {
+        self.update_archetypes(world);
+        // SAFETY: Query has unique world access.
+        unsafe {
+            self.iter_join_map_unchecked_manual(
+                world,
+                world.last_change_tick(),
+                world.read_change_tick(),
+                list,
+                map_f,
+            )
+        }
+    }
+
     /// Returns an [`Iterator`] over the query results for the given [`World`].
     ///
     /// # Safety
@@ -709,7 +753,6 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
     ///
     /// 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]
@@ -726,6 +769,29 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
         QueryManyIter::new(world, self, entities, last_change_tick, change_tick)
     }
 
+    /// Returns an [`Iterator`] over each item in an inner join on [`Entity`] between the query and a list of items which are mapped to [`Entity`]'s.
+    ///
+    /// # 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_join_map_unchecked_manual<'w, 's, I: IntoIterator, MapFn>(
+        &'s self,
+        world: &'w World,
+        last_change_tick: u32,
+        change_tick: u32,
+        list: I,
+        map_f: MapFn,
+    ) -> QueryJoinMapIter<'w, 's, Q, F, I::IntoIter, MapFn>
+    where
+        MapFn: FnMut(&I::Item) -> Entity,
+    {
+        QueryJoinMapIter::new(world, self, last_change_tick, change_tick, list, map_f)
+    }
+
     /// 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.

crates/bevy_ecs/src/system/mod.rs

@@ -134,6 +134,7 @@ mod tests {
         bundle::Bundles,
         component::{Component, Components},
         entity::{Entities, Entity},
+        event::{EventReader, Events},
         prelude::AnyOf,
         query::{Added, Changed, Or, With, Without},
         schedule::{Schedule, Stage, SystemStage},
@@ -1157,4 +1158,55 @@ mod tests {
             }
         });
     }
+
+    #[test]
+    fn iter_join_map() {
+        struct DamageEvent {
+            target: Entity,
+            damage: f32,
+        }
+
+        #[derive(Component, PartialEq, Debug)]
+        struct Health(f32);
+
+        let mut world = World::new();
+
+        let e1 = world.spawn().insert(Health(200.)).id();
+        let e2 = world.spawn().insert(Health(100.)).id();
+
+        let mut events = Events::<DamageEvent>::default();
+
+        events.extend([
+            DamageEvent {
+                target: e1,
+                damage: 50.,
+            },
+            DamageEvent {
+                target: e2,
+                damage: 80.,
+            },
+            DamageEvent {
+                target: e1,
+                damage: 150.,
+            },
+        ]);
+
+        world.insert_resource(events);
+
+        fn process_damage_events(
+            mut health_query: Query<&mut Health>,
+            mut event_reader: EventReader<DamageEvent>,
+        ) {
+            let mut join =
+                health_query.iter_join_map_mut(event_reader.iter(), |event| event.target);
+            while let Some((mut health, event)) = join.fetch_next() {
+                health.0 -= event.damage;
+            }
+        }
+        run_system(&mut world, process_damage_events);
+
+        run_system(&mut world, move |query: Query<&Health>| {
+            assert_eq!([&Health(0.), &Health(20.)], query.many([e1, e2]));
+        });
+    }
 }