Mirrors/bevy
2
0
Fork 0
mirror of https://github.com/bevyengine/bevy synced 2024-11-21 20:23:28 +00:00
Code Issues Projects Releases Packages Wiki Activity

Override QueryIter::fold to port Query::for_each perf gains to select Iterator combinators (#6773)

Browse source 
# Objective
After #6547, `Query::for_each` has been capable of automatic
vectorization on certain queries, which is seeing a notable (>50% CPU
time improvements) for iteration. However, `Query::for_each` isn't
idiomatic Rust, and lacks the flexibility of iterator combinators.

Ideally, `Query::iter` and friends should be able to achieve the same
results. However, this does seem to blocked upstream
(rust-lang/rust#104914) by Rust's loop optimizations.

## Solution
This is an intermediate solution and refactor. This moves the
`Query::for_each` implementation onto the `Iterator::fold`
implementation for `QueryIter` instead. This should result in the same
automatic vectorization optimization on all `Iterator` functions that
internally use fold, including `Iterator::for_each`, `Iterator::count`,
etc.

With this, it should close the gap between the two completely.
Internally, this PR changes `Query::for_each` to use
`query.iter().for_each(..)` instead of the duplicated implementation.

Separately, the duplicate implementations of internal iteration (i.e.
`Query::par_for_each`) now use portions of the current `Query::for_each`
implementation factored out into their own functions.

This also massively cleans up our internal fragmentation of internal
iteration options, deduplicating the iteration code used in `for_each`
and `par_iter().for_each()`.

---

## Changelog
Changed: `Query::for_each`, `Query::for_each_mut`, `Query::for_each`,
and `Query::for_each_mut` have been moved to `QueryIter`'s
`Iterator::for_each` implementation, and still retains their performance
improvements over normal iteration. These APIs are deprecated in 0.13
and will be removed in 0.14.

---------

Co-authored-by: JoJoJet <21144246+JoJoJet@users.noreply.github.com>
Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
This commit is contained in:
James Liu 2023-12-01 01:09:55 -08:00 committed by GitHub
parent e581d74f7d
commit 2148518758
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 298 additions and 199 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
benches/benches/bevy_ecs/iteration/iter_frag_foreach.rs
View file

@ -29,7 +29,7 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1.for_each_mut(&mut self.0, |mut data| {
self.1.iter_mut(&mut self.0).for_each(|mut data| {
data.0 *= 2.0;
});
}

2
benches/benches/bevy_ecs/iteration/iter_frag_foreach_sparse.rs
View file

@ -40,7 +40,7 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1.for_each_mut(&mut self.0, |mut data| {
self.1.iter_mut(&mut self.0).for_each(|mut data| {
data.0 *= 2.0;
});
}

2
benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide.rs
View file

@ -57,7 +57,7 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1.for_each_mut(&mut self.0, |mut data| {
self.1.iter_mut(&mut self.0).for_each(|mut data| {
data.0 .0 *= 2.0;
data.1 .0 *= 2.0;
data.2 .0 *= 2.0;

2
benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide_sparse.rs
View file

@ -67,7 +67,7 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1.for_each_mut(&mut self.0, |mut data| {
self.1.iter_mut(&mut self.0).for_each(|mut data| {
data.0 .0 *= 2.0;
data.1 .0 *= 2.0;
data.2 .0 *= 2.0;

3
benches/benches/bevy_ecs/iteration/iter_simple_foreach.rs
View file

@ -36,7 +36,8 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1
.for_each_mut(&mut self.0, |(velocity, mut position)| {
.iter_mut(&mut self.0)
.for_each(|(velocity, mut position)| {
position.0 += velocity.0;
});
}

3
benches/benches/bevy_ecs/iteration/iter_simple_foreach_sparse_set.rs
View file

@ -38,7 +38,8 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1
.for_each_mut(&mut self.0, |(velocity, mut position)| {
.iter_mut(&mut self.0)
.for_each(|(velocity, mut position)| {
position.0 += velocity.0;
});
}

2
benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide.rs
View file

@ -57,7 +57,7 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1.for_each_mut(&mut self.0, |mut item| {
self.1.iter_mut(&mut self.0).for_each(|mut item| {
item.1 .0 += item.0 .0;
item.3 .0 += item.2 .0;
item.5 .0 += item.4 .0;

2
benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide_sparse_set.rs
View file

@ -59,7 +59,7 @@ impl<'w> Benchmark<'w> {
#[inline(never)]
pub fn run(&mut self) {
self.1.for_each_mut(&mut self.0, |mut item| {
self.1.iter_mut(&mut self.0).for_each(|mut item| {
item.1 .0 += item.0 .0;
item.3 .0 += item.2 .0;
item.5 .0 += item.4 .0;

14
benches/benches/bevy_ecs/scheduling/running_systems.rs
View file

@ -49,17 +49,17 @@ pub fn empty_systems(criterion: &mut Criterion) {
pub fn busy_systems(criterion: &mut Criterion) {
fn ab(mut q: Query<(&mut A, &mut B)>) {
q.for_each_mut(|(mut a, mut b)| {
q.iter_mut().for_each(|(mut a, mut b)| {
std::mem::swap(&mut a.0, &mut b.0);
});
}
fn cd(mut q: Query<(&mut C, &mut D)>) {
q.for_each_mut(|(mut c, mut d)| {
q.iter_mut().for_each(|(mut c, mut d)| {
std::mem::swap(&mut c.0, &mut d.0);
});
}
fn ce(mut q: Query<(&mut C, &mut E)>) {
q.for_each_mut(|(mut c, mut e)| {
q.iter_mut().for_each(|(mut c, mut e)| {
std::mem::swap(&mut c.0, &mut e.0);
});
}
@ -98,20 +98,20 @@ pub fn busy_systems(criterion: &mut Criterion) {
pub fn contrived(criterion: &mut Criterion) {
fn s_0(mut q_0: Query<(&mut A, &mut B)>) {
q_0.for_each_mut(|(mut c_0, mut c_1)| {
q_0.iter_mut().for_each(|(mut c_0, mut c_1)| {
std::mem::swap(&mut c_0.0, &mut c_1.0);
});
}
fn s_1(mut q_0: Query<(&mut A, &mut C)>, mut q_1: Query<(&mut B, &mut D)>) {
q_0.for_each_mut(|(mut c_0, mut c_1)| {
q_0.iter_mut().for_each(|(mut c_0, mut c_1)| {
std::mem::swap(&mut c_0.0, &mut c_1.0);
});
q_1.for_each_mut(|(mut c_0, mut c_1)| {
q_1.iter_mut().for_each(|(mut c_0, mut c_1)| {
std::mem::swap(&mut c_0.0, &mut c_1.0);
});
}
fn s_2(mut q_0: Query<(&mut C, &mut D)>) {
q_0.for_each_mut(|(mut c_0, mut c_1)| {
q_0.iter_mut().for_each(|(mut c_0, mut c_1)| {
std::mem::swap(&mut c_0.0, &mut c_1.0);
});
}

6
benches/benches/bevy_ecs/scheduling/schedule.rs
View file

@ -15,19 +15,19 @@ pub fn schedule(c: &mut Criterion) {
struct E(f32);
fn ab(mut query: Query<(&mut A, &mut B)>) {
query.for_each_mut(|(mut a, mut b)| {
query.iter_mut().for_each(|(mut a, mut b)| {
std::mem::swap(&mut a.0, &mut b.0);
});
}
fn cd(mut query: Query<(&mut C, &mut D)>) {
query.for_each_mut(|(mut c, mut d)| {
query.iter_mut().for_each(|(mut c, mut d)| {
std::mem::swap(&mut c.0, &mut d.0);
});
}
fn ce(mut query: Query<(&mut C, &mut E)>) {
query.for_each_mut(|(mut c, mut e)| {
query.iter_mut().for_each(|(mut c, mut e)| {
std::mem::swap(&mut c.0, &mut e.0);
});
}

4
benches/benches/bevy_ecs/world/world_get.rs
View file

@ -230,7 +230,7 @@ pub fn world_query_for_each(criterion: &mut Criterion) {
bencher.iter(|| {
let mut count = 0;
query.for_each(&world, |comp| {
query.iter(&world).for_each(|comp| {
black_box(comp);
count += 1;
black_box(count);
@ -244,7 +244,7 @@ pub fn world_query_for_each(criterion: &mut Criterion) {
bencher.iter(|| {
let mut count = 0;
query.for_each(&world, |comp| {
query.iter(&world).for_each(|comp| {
black_box(comp);
count += 1;
black_box(count);

16
crates/bevy_ecs/src/lib.rs
View file

@ -343,7 +343,8 @@ mod tests {
let mut results = Vec::new();
world
.query::<(Entity, &A, &TableStored)>()
.for_each(&world, |(e, &i, &s)| results.push((e, i, s)));
.iter(&world)
.for_each(|(e, &i, &s)| results.push((e, i, s)));
assert_eq!(
results,
&[
@ -388,7 +389,8 @@ mod tests {
let mut results = Vec::new();
world
.query::<(Entity, &A)>()
.for_each(&world, |(e, &i)| results.push((e, i)));
.iter(&world)
.for_each(|(e, &i)| results.push((e, i)));
assert_eq!(results, &[(e, A(123)), (f, A(456))]);
}
@ -479,7 +481,8 @@ mod tests {
let mut results = Vec::new();
world
.query_filtered::<&A, With<B>>()
.for_each(&world, |i| results.push(*i));
.iter(&world)
.for_each(|i| results.push(*i));
assert_eq!(results, vec![A(123)]);
}
@ -506,7 +509,8 @@ mod tests {
let mut results = Vec::new();
world
.query_filtered::<&A, With<SparseStored>>()
.for_each(&world, |i| results.push(*i));
.iter(&world)
.for_each(|i| results.push(*i));
assert_eq!(results, vec![A(123)]);
}
@ -1395,8 +1399,8 @@ mod tests {
let mut world_a = World::new();
let world_b = World::new();
let mut query = world_a.query::<&A>();
query.for_each(&world_a, |_| {});
query.for_each(&world_b, |_| {});
query.iter(&world_a).for_each(|_| {});
query.iter(&world_b).for_each(|_| {});
}
#[test]

200
crates/bevy_ecs/src/query/iter.rs
View file

@ -1,12 +1,12 @@
use crate::{
archetype::{ArchetypeEntity, ArchetypeId, Archetypes},
archetype::{Archetype, ArchetypeEntity, ArchetypeId, Archetypes},
component::Tick,
entity::{Entities, Entity},
query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState},
storage::{TableId, TableRow, Tables},
storage::{Table, TableId, TableRow, Tables},
world::unsafe_world_cell::UnsafeWorldCell,
};
use std::{borrow::Borrow, iter::FusedIterator, mem::MaybeUninit};
use std::{borrow::Borrow, iter::FusedIterator, mem::MaybeUninit, ops::Range};
use super::{ReadOnlyWorldQueryData, WorldQueryData, WorldQueryFilter};
@ -39,6 +39,158 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> QueryIter<'w, 's, Q, F> {
cursor: QueryIterationCursor::init(world, query_state, last_run, this_run),
}
}
/// Executes the equivalent of [`Iterator::for_each`] over a contiguous segment
/// from an table.
///
/// # Safety
/// - all `rows` must be in `[0, table.entity_count)`.
/// - `table` must match Q and F
/// - Both `Q::IS_DENSE` and `F::IS_DENSE` must be true.
#[inline]
#[cfg(all(not(target = "wasm32"), feature = "multi-threaded"))]
pub(super) unsafe fn for_each_in_table_range<Func>(
&mut self,
func: &mut Func,
table: &'w Table,
rows: Range<usize>,
) where
Func: FnMut(Q::Item<'w>),
{
// SAFETY: Caller assures that Q::IS_DENSE and F::IS_DENSE are true, that table matches Q and F
// and all indicies in rows are in range.
unsafe {
self.fold_over_table_range((), &mut |_, item| func(item), table, rows);
}
}
/// Executes the equivalent of [`Iterator::for_each`] over a contiguous segment
/// from an archetype.
///
/// # Safety
/// - all `indices` must be in `[0, archetype.len())`.
/// - `archetype` must match Q and F
/// - Either `Q::IS_DENSE` or `F::IS_DENSE` must be false.
#[inline]
#[cfg(all(not(target = "wasm32"), feature = "multi-threaded"))]
pub(super) unsafe fn for_each_in_archetype_range<Func>(
&mut self,
func: &mut Func,
archetype: &'w Archetype,
rows: Range<usize>,
) where
Func: FnMut(Q::Item<'w>),
{
// SAFETY: Caller assures that either Q::IS_DENSE or F::IS_DENSE are falsae, that archetype matches Q and F
// and all indicies in rows are in range.
unsafe {
self.fold_over_archetype_range((), &mut |_, item| func(item), archetype, rows);
}
}
/// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
/// from an table.
///
/// # Safety
/// - all `rows` must be in `[0, table.entity_count)`.
/// - `table` must match Q and F
/// - Both `Q::IS_DENSE` and `F::IS_DENSE` must be true.
#[inline]
pub(super) unsafe fn fold_over_table_range<B, Func>(
&mut self,
mut accum: B,
func: &mut Func,
table: &'w Table,
rows: Range<usize>,
) -> B
where
Func: FnMut(B, Q::Item<'w>) -> B,
{
Q::set_table(&mut self.cursor.fetch, &self.query_state.fetch_state, table);
F::set_table(
&mut self.cursor.filter,
&self.query_state.filter_state,
table,
);
let entities = table.entities();
for row in rows {
// SAFETY: Caller assures `row` in range of the current archetype.
let entity = entities.get_unchecked(row);
let row = TableRow::new(row);
// SAFETY: set_table was called prior.
// Caller assures `row` in range of the current archetype.
if !F::filter_fetch(&mut self.cursor.filter, *entity, row) {
continue;
}
// SAFETY: set_table was called prior.
// Caller assures `row` in range of the current archetype.
let item = Q::fetch(&mut self.cursor.fetch, *entity, row);
accum = func(accum, item);
}
accum
}
/// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
/// from an archetype.
///
/// # Safety
/// - all `indices` must be in `[0, archetype.len())`.
/// - `archetype` must match Q and F
/// - Either `Q::IS_DENSE` or `F::IS_DENSE` must be false.
#[inline]
pub(super) unsafe fn fold_over_archetype_range<B, Func>(
&mut self,
mut accum: B,
func: &mut Func,
archetype: &'w Archetype,
indices: Range<usize>,
) -> B
where
Func: FnMut(B, Q::Item<'w>) -> B,
{
let table = self.tables.get(archetype.table_id()).debug_checked_unwrap();
Q::set_archetype(
&mut self.cursor.fetch,
&self.query_state.fetch_state,
archetype,
table,
);
F::set_archetype(
&mut self.cursor.filter,
&self.query_state.filter_state,
archetype,
table,
);
let entities = archetype.entities();
for index in indices {
// SAFETY: Caller assures `index` in range of the current archetype.
let archetype_entity = entities.get_unchecked(index);
// SAFETY: set_archetype was called prior.
// Caller assures `index` in range of the current archetype.
if !F::filter_fetch(
&mut self.cursor.filter,
archetype_entity.entity(),
archetype_entity.table_row(),
) {
continue;
}
// SAFETY: set_archetype was called prior, `index` is an archetype index in range of the current archetype
// Caller assures `index` in range of the current archetype.
let item = Q::fetch(
&mut self.cursor.fetch,
archetype_entity.entity(),
archetype_entity.table_row(),
);
accum = func(accum, item);
}
accum
}
}
impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> Iterator for QueryIter<'w, 's, Q, F> {
@ -61,6 +213,44 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> Iterator for QueryIter<'w,
let min_size = if archetype_query { max_size } else { 0 };
(min_size, Some(max_size))
}
#[inline]
fn fold<B, Func>(mut self, init: B, mut func: Func) -> B
where
Func: FnMut(B, Self::Item) -> B,
{
let mut accum = init;
// Empty any remaining uniterated values from the current table/archetype
while self.cursor.current_row != self.cursor.current_len {
let Some(item) = self.next() else { break };
accum = func(accum, item);
}
if Q::IS_DENSE && F::IS_DENSE {
for table_id in self.cursor.table_id_iter.clone() {
// SAFETY: Matched table IDs are guaranteed to still exist.
let table = unsafe { self.tables.get(*table_id).debug_checked_unwrap() };
accum =
// SAFETY:
// - The fetched table matches both Q and F
// - The provided range is equivalent to [0, table.entity_count)
// - The if block ensures that Q::IS_DENSE and F::IS_DENSE are both true
unsafe { self.fold_over_table_range(accum, &mut func, table, 0..table.entity_count()) };
}
} else {
for archetype_id in self.cursor.archetype_id_iter.clone() {
let archetype =
// SAFETY: Matched archetype IDs are guaranteed to still exist.
unsafe { self.archetypes.get(*archetype_id).debug_checked_unwrap() };
accum =
// SAFETY:
// - The fetched archetype matches both Q and F
// - The provided range is equivalent to [0, archetype.len)
// - The if block ensures that ether Q::IS_DENSE or F::IS_DENSE are false
unsafe { self.fold_over_archetype_range(accum, &mut func, archetype, 0..archetype.len()) };
}
}
accum
}
}
// This is correct as [`QueryIter`] always returns `None` once exhausted.
@ -547,7 +737,7 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> QueryIterationCursor<'w, 's
}
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
// QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::par_for_each_unchecked_manual
/// # Safety
/// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`]
/// was initialized for.
@ -599,9 +789,9 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> QueryIterationCursor<'w, 's
if self.current_row == self.current_len {
let archetype_id = self.archetype_id_iter.next()?;
let archetype = archetypes.get(*archetype_id).debug_checked_unwrap();
let table = tables.get(archetype.table_id()).debug_checked_unwrap();
// SAFETY: `archetype` and `tables` are from the world that `fetch/filter` were created for,
// `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
let table = tables.get(archetype.table_id()).debug_checked_unwrap();
Q::set_archetype(&mut self.fetch, &query_state.fetch_state, archetype, table);
F::set_archetype(
&mut self.filter,

4
crates/bevy_ecs/src/query/mod.rs
View file

@ -751,7 +751,7 @@ mod tests {
let _: Option<[&Foo; 2]> = q.iter_combinations::<2>(&world).next();
let _: Option<&Foo> = q.iter_manual(&world).next();
let _: Option<&Foo> = q.iter_many(&world, [e]).next();
q.for_each(&world, |_: &Foo| ());
q.iter(&world).for_each(|_: &Foo| ());
let _: Option<&Foo> = q.get(&world, e).ok();
let _: Option<&Foo> = q.get_manual(&world, e).ok();
@ -765,7 +765,7 @@ mod tests {
let _: Option<&Foo> = q.iter().next();
let _: Option<[&Foo; 2]> = q.iter_combinations::<2>().next();
let _: Option<&Foo> = q.iter_many([e]).next();
q.for_each(|_: &Foo| ());
q.iter().for_each(|_: &Foo| ());
let _: Option<&Foo> = q.get(e).ok();
let _: Option<&Foo> = q.get_component(e).ok();

18
crates/bevy_ecs/src/query/par_iter.rs
View file

@ -118,12 +118,9 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> QueryParIter<'w, 's, Q, F>
// Query or a World, which ensures that multiple aliasing QueryParIters cannot exist
// at the same time.
unsafe {
self.state.for_each_unchecked_manual(
self.world,
func,
self.last_run,
self.this_run,
);
self.state
.iter_unchecked_manual(self.world, self.last_run, self.this_run)
.for_each(func);
}
}
#[cfg(all(not(target = "wasm32"), feature = "multi-threaded"))]
@ -132,12 +129,9 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> QueryParIter<'w, 's, Q, F>
if thread_count <= 1 {
// SAFETY: See the safety comment above.
unsafe {
self.state.for_each_unchecked_manual(
self.world,
func,
self.last_run,
self.this_run,
);
self.state
.iter_unchecked_manual(self.world, self.last_run, self.this_run)
.for_each(func);
}
} else {
// Need a batch size of at least 1.

166
crates/bevy_ecs/src/query/state.rs
View file

@ -8,7 +8,7 @@ use crate::{
Access, BatchingStrategy, DebugCheckedUnwrap, FilteredAccess, QueryCombinationIter,
QueryIter, QueryParIter,
},
storage::{TableId, TableRow},
storage::TableId,
world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
};
#[cfg(feature = "trace")]
@ -1008,36 +1008,28 @@ impl<Q: WorldQueryData, F: WorldQueryFilter> QueryState<Q, F> {
/// iter() method, but cannot be chained like a normal [`Iterator`].
///
/// This can only be called for read-only queries, see [`Self::for_each_mut`] for write-queries.
///
/// Shorthand for `query.iter(world).for_each(..)`.
#[inline]
#[deprecated(
since = "0.13.0",
note = "QueryState::for_each was not idiomatic Rust and has been moved to query.iter().for_each()"
)]
pub fn for_each<'w, FN: FnMut(ROQueryItem<'w, Q>)>(&mut self, world: &'w World, func: FN) {
self.update_archetypes(world);
// SAFETY: query is read only
unsafe {
self.as_readonly().for_each_unchecked_manual(
world.as_unsafe_world_cell_readonly(),
func,
world.last_change_tick(),
world.read_change_tick(),
);
}
self.iter(world).for_each(func);
}
/// Runs `func` on each query result for the given [`World`]. This is faster than the equivalent
/// `iter_mut()` method, but cannot be chained like a normal [`Iterator`].
///
/// Shorthand for `query.iter_mut(world).for_each(..)`.
#[inline]
#[deprecated(
since = "0.13.0",
note = "QueryState::for_each_mut was not idiomatic Rust and has been moved to query.iter_mut().for_each()"
)]
pub fn for_each_mut<'w, FN: FnMut(Q::Item<'w>)>(&mut self, world: &'w mut World, func: FN) {
self.update_archetypes(world);
let change_tick = world.change_tick();
let last_change_tick = world.last_change_tick();
// SAFETY: query has unique world access
unsafe {
self.for_each_unchecked_manual(
world.as_unsafe_world_cell(),
func,
last_change_tick,
change_tick,
);
}
self.iter_mut(world).for_each(func);
}
/// Runs `func` on each query result for the given [`World`]. This is faster than the equivalent
@ -1054,7 +1046,8 @@ impl<Q: WorldQueryData, F: WorldQueryFilter> QueryState<Q, F> {
func: FN,
) {
self.update_archetypes_unsafe_world_cell(world);
self.for_each_unchecked_manual(world, func, world.last_change_tick(), world.change_tick());
self.iter_unchecked_manual(world, world.last_change_tick(), world.change_tick())
.for_each(func);
}
/// Returns a parallel iterator over the query results for the given [`World`].
@ -1096,73 +1089,6 @@ impl<Q: WorldQueryData, F: WorldQueryFilter> QueryState<Q, F> {
}
}
/// 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`].
///
/// # 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 for_each_unchecked_manual<'w, FN: FnMut(Q::Item<'w>)>(
&self,
world: UnsafeWorldCell<'w>,
mut func: FN,
last_run: Tick,
this_run: Tick,
) {
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
let mut fetch = Q::init_fetch(world, &self.fetch_state, last_run, this_run);
let mut filter = F::init_fetch(world, &self.filter_state, last_run, this_run);
let tables = &world.storages().tables;
if Q::IS_DENSE && F::IS_DENSE {
for table_id in &self.matched_table_ids {
let table = tables.get(*table_id).debug_checked_unwrap();
Q::set_table(&mut fetch, &self.fetch_state, table);
F::set_table(&mut filter, &self.filter_state, table);
let entities = table.entities();
for row in 0..table.entity_count() {
let entity = entities.get_unchecked(row);
let row = TableRow::new(row);
if !F::filter_fetch(&mut filter, *entity, row) {
continue;
}
func(Q::fetch(&mut fetch, *entity, row));
}
}
} else {
let archetypes = world.archetypes();
for archetype_id in &self.matched_archetype_ids {
let archetype = archetypes.get(*archetype_id).debug_checked_unwrap();
let table = tables.get(archetype.table_id()).debug_checked_unwrap();
Q::set_archetype(&mut fetch, &self.fetch_state, archetype, table);
F::set_archetype(&mut filter, &self.filter_state, archetype, table);
let entities = archetype.entities();
for idx in 0..archetype.len() {
let archetype_entity = entities.get_unchecked(idx);
if !F::filter_fetch(
&mut filter,
archetype_entity.entity(),
archetype_entity.table_row(),
) {
continue;
}
func(Q::fetch(
&mut fetch,
archetype_entity.entity(),
archetype_entity.table_row(),
));
}
}
}
}
/// Runs `func` on each query result in parallel 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`].
@ -1205,28 +1131,19 @@ impl<Q: WorldQueryData, F: WorldQueryFilter> QueryState<Q, F> {
let mut offset = 0;
while offset < table.entity_count() {
let func = func.clone();
let mut func = func.clone();
let len = batch_size.min(table.entity_count() - offset);
scope.spawn(async move {
#[cfg(feature = "trace")]
let _span = self.par_iter_span.enter();
let mut fetch =
Q::init_fetch(world, &self.fetch_state, last_run, this_run);
let mut filter =
F::init_fetch(world, &self.filter_state, last_run, this_run);
let tables = &world.storages().tables;
let table = tables.get(*table_id).debug_checked_unwrap();
let entities = table.entities();
Q::set_table(&mut fetch, &self.fetch_state, table);
F::set_table(&mut filter, &self.filter_state, table);
for row in offset..offset + len {
let entity = entities.get_unchecked(row);
let row = TableRow::new(row);
if !F::filter_fetch(&mut filter, *entity, row) {
continue;
}
func(Q::fetch(&mut fetch, *entity, row));
}
let table = &world
.storages()
.tables
.get(*table_id)
.debug_checked_unwrap();
let batch = offset..offset + len;
self.iter_unchecked_manual(world, last_run, this_run)
.for_each_in_table_range(&mut func, table, batch);
});
offset += batch_size;
}
@ -1241,40 +1158,17 @@ impl<Q: WorldQueryData, F: WorldQueryFilter> QueryState<Q, F> {
}
while offset < archetype.len() {
let func = func.clone();
let mut func = func.clone();
let len = batch_size.min(archetype.len() - offset);
scope.spawn(async move {
#[cfg(feature = "trace")]
let _span = self.par_iter_span.enter();
let mut fetch =
Q::init_fetch(world, &self.fetch_state, last_run, this_run);
let mut filter =
F::init_fetch(world, &self.filter_state, last_run, this_run);
let tables = &world.storages().tables;
let archetype =
world.archetypes().get(*archetype_id).debug_checked_unwrap();
let table = tables.get(archetype.table_id()).debug_checked_unwrap();
Q::set_archetype(&mut fetch, &self.fetch_state, archetype, table);
F::set_archetype(&mut filter, &self.filter_state, archetype, table);
let entities = archetype.entities();
for archetype_row in offset..offset + len {
let archetype_entity = entities.get_unchecked(archetype_row);
if !F::filter_fetch(
&mut filter,
archetype_entity.entity(),
archetype_entity.table_row(),
) {
continue;
}
func(Q::fetch(
&mut fetch,
archetype_entity.entity(),
archetype_entity.table_row(),
));
}
let batch = offset..offset + len;
self.iter_unchecked_manual(world, last_run, this_run)
.for_each_in_archetype_range(&mut func, archetype, batch);
});
offset += batch_size;
}
}

25
crates/bevy_ecs/src/system/query.rs
View file

@ -712,6 +712,8 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> Query<'w, 's, Q, F> {
/// Runs `f` on each read-only query item.
///
/// Shorthand for `query.iter().for_each(..)`.
///
/// # Example
///
/// Here, the `report_names_system` iterates over the `Player` component of every entity that contains it:
@ -735,22 +737,26 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> Query<'w, 's, Q, F> {
/// - [`for_each_mut`](Self::for_each_mut) to operate on mutable query items.
/// - [`iter`](Self::iter) for the iterator based alternative.
#[inline]
#[deprecated(
since = "0.13.0",
note = "Query::for_each was not idiomatic Rust and has been moved to query.iter().for_each()"
)]
pub fn for_each<'this>(&'this self, f: impl FnMut(ROQueryItem<'this, Q>)) {
// SAFETY:
// - `self.world` has permission to access the required components.
// - The query is read-only, so it can be aliased even if it was originally mutable.
unsafe {
self.state.as_readonly().for_each_unchecked_manual(
self.world,
f,
self.last_run,
self.this_run,
);
self.state
.as_readonly()
.iter_unchecked_manual(self.world, self.last_run, self.this_run)
.for_each(f);
};
}
/// Runs `f` on each query item.
///
/// Shorthand for `query.iter_mut().for_each(..)`.
///
/// # Example
///
/// Here, the `gravity_system` updates the `Velocity` component of every entity that contains it:
@ -774,11 +780,16 @@ impl<'w, 's, Q: WorldQueryData, F: WorldQueryFilter> Query<'w, 's, Q, F> {
/// - [`for_each`](Self::for_each) to operate on read-only query items.
/// - [`iter_mut`](Self::iter_mut) for the iterator based alternative.
#[inline]
#[deprecated(
since = "0.13.0",
note = "Query::for_each_mut was not idiomatic Rust and has been moved to query.iter_mut().for_each()"
)]
pub fn for_each_mut<'a>(&'a mut self, f: impl FnMut(Q::Item<'a>)) {
// SAFETY: `self.world` has permission to access the required components.
unsafe {
self.state
.for_each_unchecked_manual(self.world, f, self.last_run, self.this_run);
.iter_unchecked_manual(self.world, self.last_run, self.this_run)
.for_each(f);
};
}

8
crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs
View file

@ -75,16 +75,16 @@ fn main() {
{
let mut opt_data: Option<&Foo> = None;
let mut opt_data_2: Option<Mut<Foo>> = None;
query.for_each(|data| opt_data = Some(data));
query.for_each_mut(|data| opt_data_2 = Some(data));
query.iter().for_each(|data| opt_data = Some(data));
query.iter_mut().for_each(|data| opt_data_2 = Some(data));
assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
}
{
let mut opt_data_2: Option<Mut<Foo>> = None;
let mut opt_data: Option<&Foo> = None;
query.for_each_mut(|data| opt_data_2 = Some(data));
query.for_each(|data| opt_data = Some(data));
query.iter_mut().for_each(|data| opt_data_2 = Some(data));
query.iter().for_each(|data| opt_data = Some(data));
assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
}
dbg!("bye");

10
crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr
View file

@ -101,19 +101,19 @@ error[E0502]: cannot borrow `query` as immutable because it is also borrowed as
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:79:13
|
78 | query.for_each(|data| opt_data = Some(data));
78 | query.iter().for_each(|data| opt_data = Some(data));
| ----- immutable borrow occurs here
79 | query.for_each_mut(|data| opt_data_2 = Some(data));
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
79 | query.iter_mut().for_each(|data| opt_data_2 = Some(data));
| ^^^^^^^^^^^^^^^^ mutable borrow occurs here
80 | assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
| -------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:87:13
|
86 | query.for_each_mut(|data| opt_data_2 = Some(data));
86 | query.iter_mut().for_each(|data| opt_data_2 = Some(data));
| ----- mutable borrow occurs here
87 | query.for_each(|data| opt_data = Some(data));
87 | query.iter().for_each(|data| opt_data = Some(data));
| ^^^^^ immutable borrow occurs here
88 | assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
| ---------- mutable borrow later used here

8
examples/stress_tests/many_buttons.rs
View file

@ -67,13 +67,17 @@ fn main() {
if args.relayout {
app.add_systems(Update, |mut style_query: Query<&mut Style>| {
style_query.for_each_mut(|mut style| style.set_changed());
style_query
.iter_mut()
.for_each(|mut style| style.set_changed());
});
}
if args.recompute_text {
app.add_systems(Update, |mut text_query: Query<&mut Text>| {
text_query.for_each_mut(|mut text| text.set_changed());
text_query
.iter_mut()
.for_each(|mut text| text.set_changed());
});
}