Replace many_for_each_mut with iter_many_mut. (#5402)

# Objective
Replace `many_for_each_mut` with `iter_many_mut` using the same tricks to avoid aliased mutability that `iter_combinations_mut` uses.

<sub>I tried rebasing the draft PR I made for this before and it died. F</sub>
## Why
`many_for_each_mut` is worse for a few reasons:
1. The closure prevents the use of `continue`, `break`, and `return` behaves like a limited `continue`.
2. rustfmt will crumple it and double the indentation when the line gets too long.
    ```rust
    query.many_for_each_mut(
        &entity_list,
        |(mut transform, velocity, mut component_c)| {
            // Double trouble.
        },
    );
    ```
3. It is more surprising to have `many_for_each_mut` as a mutable counterpart to `iter_many` than `iter_many_mut`.
4. It required a separate unsafe fn; more unsafe code to maintain.
5. The `iter_many_mut` API matches the existing `iter_combinations_mut` API.

Co-authored-by: devil-ira <justthecooldude@gmail.com>

Author: tim-blackbird

crates/bevy_ecs/src/query/iter.rs

@@ -72,9 +72,9 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Iterator for QueryIter<'w, 's, Q, F>
 // This is correct as [`QueryIter`] always returns `None` once exhausted.
 impl<'w, 's, Q: WorldQuery, F: WorldQuery> FusedIterator for QueryIter<'w, 's, Q, F> {}
 
-/// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
+/// An [`Iterator`] over [`Query`](crate::system::Query) results of a list of [`Entity`]s.
 ///
-/// This struct is created by the [`Query::iter_many`](crate::system::Query::iter_many) method.
+/// This struct is created by the [`Query::iter_many`](crate::system::Query::iter_many) and [`Query::iter_many_mut`](crate::system::Query::iter_many_mut) methods.
 pub struct QueryManyIter<'w, 's, Q: WorldQuery, F: WorldQuery, I: Iterator>
 where
     I::Item: Borrow<Entity>,
@@ -126,16 +126,15 @@ where
             entity_iter: entity_list.into_iter(),
         }
     }
-}
-
-impl<'w, 's, Q: WorldQuery, F: WorldQuery, I: Iterator> Iterator for QueryManyIter<'w, 's, Q, F, I>
-where
-    I::Item: Borrow<Entity>,
-{
-    type Item = QueryItem<'w, Q>;
 
+    /// 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 shared access.
     #[inline(always)]
-    fn next(&mut self) -> Option<Self::Item> {
+    unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<QueryItem<'w, Q>> {
         for entity in self.entity_iter.by_ref() {
             let location = match self.entities.get(*entity.borrow()) {
                 Some(location) => location,
@@ -154,32 +153,61 @@ where
 
             // SAFETY: `archetype` is from the world that `fetch/filter` were created for,
             // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
-            unsafe {
-                self.fetch
-                    .set_archetype(&self.query_state.fetch_state, archetype, self.tables);
-            }
+            self.fetch
+                .set_archetype(&self.query_state.fetch_state, archetype, self.tables);
+
             // SAFETY: `table` is from the world that `fetch/filter` were created for,
             // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
-            unsafe {
-                self.filter
-                    .set_archetype(&self.query_state.filter_state, archetype, self.tables);
-            }
+            self.filter
+                .set_archetype(&self.query_state.filter_state, archetype, self.tables);
+
             // SAFETY: set_archetype was called prior.
             // `location.index` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d
-            if unsafe { self.filter.archetype_filter_fetch(location.index) } {
+            if self.filter.archetype_filter_fetch(location.index) {
                 // SAFETY: set_archetype was called prior, `location.index` is an archetype index in range of the current archetype
-                return Some(unsafe { self.fetch.archetype_fetch(location.index) });
+                return Some(self.fetch.archetype_fetch(location.index));
             }
         }
         None
     }
 
+    /// Get next result from the query
+    #[inline(always)]
+    pub fn fetch_next(&mut self) -> Option<QueryItem<'_, Q>> {
+        // 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::shrink) }
+    }
+}
+
+impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator> Iterator
+    for QueryManyIter<'w, 's, Q, F, I>
+where
+    I::Item: Borrow<Entity>,
+{
+    type Item = QueryItem<'w, Q>;
+
+    #[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.entity_iter.size_hint();
         (0, max_size)
     }
 }
 
+// This is correct as [`QueryManyIter`] always returns `None` once exhausted.
+impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator> FusedIterator
+    for QueryManyIter<'w, 's, Q, F, I>
+where
+    I::Item: Borrow<Entity>,
+{
+}
+
 pub struct QueryCombinationIter<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> {
     tables: &'w Tables,
     archetypes: &'w Archetypes,
@@ -476,7 +504,7 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> QueryIterationCursor<'w, 's, Q, F> {
     }
 
     // NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
-    // QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
+    // QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
     /// # Safety
     /// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`]
     /// was initialized for.

crates/bevy_ecs/src/query/mod.rs

@@ -633,9 +633,10 @@ count():       {count}"#
         }
         {
             fn system(has_a: Query<Entity, With<A>>, mut b_query: Query<&mut B>) {
-                b_query.many_for_each_mut(&has_a, |mut b| {
+                let mut iter = b_query.iter_many_mut(&has_a);
+                while let Some(mut b) = iter.fetch_next() {
                     b.0 = 1;
-                });
+                }
             }
             let mut system = IntoSystem::into_system(system);
             system.initialize(&mut world);

crates/bevy_ecs/src/query/state.rs

@@ -602,7 +602,7 @@ 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.
+    /// See [`Self::iter_many_mut`] for queries that contain at least one mutable component.
     ///
     #[inline]
     pub fn iter_many<'w, 's, EntityList: IntoIterator>(
@@ -613,9 +613,9 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
     where
         EntityList::Item: Borrow<Entity>,
     {
+        self.update_archetypes(world);
         // SAFETY: query is read only
         unsafe {
-            self.update_archetypes(world);
             self.as_readonly().iter_many_unchecked_manual(
                 entities,
                 world,
@@ -625,6 +625,28 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
         }
     }
 
+    /// Returns an iterator over the query results of a list of [`Entity`]'s.
+    #[inline]
+    pub fn iter_many_mut<'w, 's, EntityList: IntoIterator>(
+        &'s mut self,
+        world: &'w mut World,
+        entities: EntityList,
+    ) -> QueryManyIter<'w, 's, Q, F, EntityList::IntoIter>
+    where
+        EntityList::Item: Borrow<Entity>,
+    {
+        self.update_archetypes(world);
+        // SAFETY: Query has unique world access.
+        unsafe {
+            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
@@ -868,29 +890,6 @@ 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`].
@@ -909,7 +908,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::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
+        // QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::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(
@@ -978,7 +977,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::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
+        // QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
         ComputeTaskPool::get().scope(|scope| {
             if <QueryFetch<'static, Q>>::IS_DENSE && <QueryFetch<'static, F>>::IS_DENSE {
                 let tables = &world.storages().tables;
@@ -1078,62 +1077,6 @@ 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 {
-            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.
     ///