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>
This commit is contained in:
ira 2022-07-30 01:38:13 +00:00
parent 418beffaec
commit 83a9e16158
11 changed files with 217 additions and 169 deletions

View file

@ -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.

View file

@ -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);

View file

@ -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.
///

View file

@ -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

View file

@ -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() {}

View file

@ -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`

View file

@ -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() {}

View file

@ -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`

View file

@ -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() {}

View file

@ -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

View file

@ -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