Skip to content

Commit

Permalink
Replace many_for_each_mut with iter_many_mut. (bevyengine#5402)
Browse files Browse the repository at this point in the history 
# 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>
  • Loading branch information
@tim-blackbird @inodentry
2 people authored and inodentry committed Aug 8, 2022
1 parent 9e7cfbd commit fe04766
Show file tree
Hide file tree
Showing 11 changed files with 217 additions and 169 deletions.
70 changes: 49 additions & 21 deletions crates/bevy_ecs/src/query/iter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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.
Expand Down
5 changes: 3 additions & 2 deletions crates/bevy_ecs/src/query/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
109 changes: 26 additions & 83 deletions crates/bevy_ecs/src/query/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>(
Expand All @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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`].
Expand All @@ -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(
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
///
Expand Down
Loading

0 comments on commit fe04766

Please sign in to comment.