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>

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
-                self.fetch
+                .set_archetype(&self.query_state.fetch_state, archetype, self.tables);
-                    .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
-                self.filter
+                .set_archetype(&self.query_state.filter_state, archetype, self.tables);
-                    .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.
     ///

crates/bevy_ecs/src/system/query.rs

@@ -424,7 +424,7 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, 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.
     ///
     /// # Examples
     /// ```
@@ -471,6 +471,55 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
         }
     }
 
+    /// Calls a closure on each result of [`Query`] where the entities match.
+    /// # 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 {
+    ///         let mut iter = counter_query.iter_many_mut(&friends.list);
+    ///         while let Some(mut counter) = iter.fetch_next() {
+    ///             println!("Friend's counter: {:?}", counter.value);
+    ///             counter.value += 1;
+    ///         }
+    ///     }
+    /// }
+    /// # bevy_ecs::system::assert_is_system(system);
+    /// ```
+    #[inline]
+    pub fn iter_many_mut<EntityList: IntoIterator>(
+        &mut self,
+        entities: EntityList,
+    ) -> QueryManyIter<'_, '_, 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.
     ///
     /// # Safety
@@ -506,7 +555,7 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
 
     /// Returns an [`Iterator`] over the query results of a list of [`Entity`]'s.
     ///
-    /// If you want safe mutable access to query results of a list of [`Entity`]'s. See [`Self::many_for_each_mut`].
+    /// If you want safe mutable access to query results of a list of [`Entity`]'s. See [`Self::iter_many_mut`].
     ///
     /// # Safety
     /// This allows aliased mutability and does not check for entity uniqueness.
@@ -670,55 +719,6 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
         };
     }
 
-    /// Calls a closure on each result of [`Query`] where the entities match.
-    /// # 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 {
-    ///         counter_query.many_for_each_mut(&friends.list, |mut counter| {
-    ///             println!("Friend's counter: {:?}", counter.value);
-    ///             counter.value += 1;
-    ///         });
-    ///     }
-    /// }
-    /// # bevy_ecs::system::assert_is_system(system);
-    /// ```
-    #[inline]
-    pub fn many_for_each_mut<EntityList: IntoIterator>(
-        &mut self,
-        entities: EntityList,
-        f: impl FnMut(QueryItem<'_, Q>),
-    ) 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.many_for_each_unchecked_manual(
-                self.world,
-                entities,
-                f,
-                self.last_change_tick,
-                self.change_tick,
-            );
-        };
-    }
-
     /// Returns the query result for the given [`Entity`].
     ///
     /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is

crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.rs

@@ -0,0 +1,15 @@
+use bevy_ecs::prelude::*;
+
+#[derive(Component)]
+struct A(usize);
+
+fn system(mut query: Query<&mut A>) {
+    let iter = query.iter_combinations_mut();
+
+    // This should fail to compile.
+    is_iterator(iter)
+}
+
+fn is_iterator<T: Iterator>(_iter: T) {}
+
+fn main() {}

crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.stderr

@@ -0,0 +1,25 @@
+error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied
+  --> tests/ui/query_iter_combinations_mut_iterator_safety.rs:10:17
+   |
+10 |     is_iterator(iter)
+   |     ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A`
+   |     |
+   |     required by a bound introduced by this call
+   |
+   = help: the following other types implement trait `ReadOnlyWorldQuery`:
+             &T
+             ()
+             (F0, F1)
+             (F0, F1, F2)
+             (F0, F1, F2, F3)
+             (F0, F1, F2, F3, F4)
+             (F0, F1, F2, F3, F4, F5)
+             (F0, F1, F2, F3, F4, F5, F6)
+           and 49 others
+   = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A`
+   = note: required because of the requirements on the impl of `Iterator` for `QueryCombinationIter<'_, '_, &mut A, (), {_: usize}>`
+note: required by a bound in `is_iterator`
+  --> tests/ui/query_iter_combinations_mut_iterator_safety.rs:13:19
+   |
+13 | fn is_iterator<T: Iterator>(_iter: T) {}
+   |                   ^^^^^^^^ required by this bound in `is_iterator`

crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.rs

@@ -0,0 +1,15 @@
+use bevy_ecs::prelude::*;
+
+#[derive(Component)]
+struct A(usize);
+
+fn system(mut query: Query<&mut A>, e: Res<Entity>) {
+    let iter = query.iter_many_mut([*e]);
+
+    // This should fail to compile.
+    is_iterator(iter)
+}
+
+fn is_iterator<T: Iterator>(_iter: T) {}
+
+fn main() {}

crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.stderr

@@ -0,0 +1,25 @@
+error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied
+  --> tests/ui/query_iter_many_mut_iterator_safety.rs:10:17
+   |
+10 |     is_iterator(iter)
+   |     ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A`
+   |     |
+   |     required by a bound introduced by this call
+   |
+   = help: the following other types implement trait `ReadOnlyWorldQuery`:
+             &T
+             ()
+             (F0, F1)
+             (F0, F1, F2)
+             (F0, F1, F2, F3)
+             (F0, F1, F2, F3, F4)
+             (F0, F1, F2, F3, F4, F5)
+             (F0, F1, F2, F3, F4, F5, F6)
+           and 49 others
+   = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A`
+   = note: required because of the requirements on the impl of `Iterator` for `QueryManyIter<'_, '_, &mut A, (), std::array::IntoIter<bevy_ecs::entity::Entity, 1_usize>>`
+note: required by a bound in `is_iterator`
+  --> tests/ui/query_iter_many_mut_iterator_safety.rs:13:19
+   |
+13 | fn is_iterator<T: Iterator>(_iter: T) {}
+   |                   ^^^^^^^^ required by this bound in `is_iterator`

crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_many_mut_lifetime_safety.rs

@@ -5,10 +5,11 @@ struct A(usize);
 
 fn system(mut query: Query<&mut A>, e: Res<Entity>) {
     let mut results = Vec::new();
-    query.many_for_each_mut(vec![*e, *e], |a| {
+    let mut iter = query.iter_many_mut([*e, *e]);
+    while let Some(a) = iter.fetch_next() {
         // this should fail to compile
         results.push(a);
-    });
+    }
 }
 
 fn main() {}

crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_many_mut_lifetime_safety.stderr

@@ -0,0 +1,5 @@
+error[E0499]: cannot borrow `iter` as mutable more than once at a time
+ --> tests/ui/system_query_iter_many_mut_lifetime_safety.rs:9:25
+  |
+9 |     while let Some(a) = iter.fetch_next() {
+  |                         ^^^^^^^^^^^^^^^^^ `iter` was mutably borrowed here in the previous iteration of the loop

crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_many_for_each_mut_lifetime_safety.stderr

@@ -1,10 +0,0 @@
-error[E0521]: borrowed data escapes outside of closure
-  --> tests/ui/system_query_many_for_each_mut_lifetime_safety.rs:10:9
-   |
-7  |     let mut results = Vec::new();
-   |         ----------- `results` declared here, outside of the closure body
-8  |     query.many_for_each_mut(vec![*e, *e], |a| {
-   |                                            - `a` is a reference that is only valid in the closure body
-9  |         // this should fail to compile
-10 |         results.push(a);
-   |         ^^^^^^^^^^^^^^^ `a` escapes the closure body here