From 2148518758aaae30a2f04a5758cda5bf7c9154ca Mon Sep 17 00:00:00 2001 From: James Liu Date: Fri, 1 Dec 2023 01:09:55 -0800 Subject: [PATCH] Override QueryIter::fold to port Query::for_each perf gains to select Iterator combinators (#6773) # 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 --- .../bevy_ecs/iteration/iter_frag_foreach.rs | 2 +- .../iteration/iter_frag_foreach_sparse.rs | 2 +- .../iteration/iter_frag_foreach_wide.rs | 2 +- .../iter_frag_foreach_wide_sparse.rs | 2 +- .../bevy_ecs/iteration/iter_simple_foreach.rs | 3 +- .../iter_simple_foreach_sparse_set.rs | 3 +- .../iteration/iter_simple_foreach_wide.rs | 2 +- .../iter_simple_foreach_wide_sparse_set.rs | 2 +- .../bevy_ecs/scheduling/running_systems.rs | 14 +- .../benches/bevy_ecs/scheduling/schedule.rs | 6 +- benches/benches/bevy_ecs/world/world_get.rs | 4 +- crates/bevy_ecs/src/lib.rs | 16 +- crates/bevy_ecs/src/query/iter.rs | 200 +++++++++++++++++- crates/bevy_ecs/src/query/mod.rs | 4 +- crates/bevy_ecs/src/query/par_iter.rs | 18 +- crates/bevy_ecs/src/query/state.rs | 166 +++------------ crates/bevy_ecs/src/system/query.rs | 25 ++- .../tests/ui/query_lifetime_safety.rs | 8 +- .../tests/ui/query_lifetime_safety.stderr | 10 +- examples/stress_tests/many_buttons.rs | 8 +- 20 files changed, 298 insertions(+), 199 deletions(-) diff --git a/benches/benches/bevy_ecs/iteration/iter_frag_foreach.rs b/benches/benches/bevy_ecs/iteration/iter_frag_foreach.rs index b4a8877074e44..c345fe9008749 100644 --- a/benches/benches/bevy_ecs/iteration/iter_frag_foreach.rs +++ b/benches/benches/bevy_ecs/iteration/iter_frag_foreach.rs @@ -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; }); } diff --git a/benches/benches/bevy_ecs/iteration/iter_frag_foreach_sparse.rs b/benches/benches/bevy_ecs/iteration/iter_frag_foreach_sparse.rs index 0d92b39999724..c616f737c21a8 100644 --- a/benches/benches/bevy_ecs/iteration/iter_frag_foreach_sparse.rs +++ b/benches/benches/bevy_ecs/iteration/iter_frag_foreach_sparse.rs @@ -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; }); } diff --git a/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide.rs b/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide.rs index f385b04ebd6d8..0ea09053b8251 100644 --- a/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide.rs +++ b/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide.rs @@ -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; diff --git a/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide_sparse.rs b/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide_sparse.rs index e21887204d26b..6949d1b90d23b 100644 --- a/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide_sparse.rs +++ b/benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide_sparse.rs @@ -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; diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_foreach.rs b/benches/benches/bevy_ecs/iteration/iter_simple_foreach.rs index 05956851e50c7..6dae375ed7a43 100644 --- a/benches/benches/bevy_ecs/iteration/iter_simple_foreach.rs +++ b/benches/benches/bevy_ecs/iteration/iter_simple_foreach.rs @@ -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; }); } diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_foreach_sparse_set.rs b/benches/benches/bevy_ecs/iteration/iter_simple_foreach_sparse_set.rs index f9396f0ff287f..8a9e7baaab6c2 100644 --- a/benches/benches/bevy_ecs/iteration/iter_simple_foreach_sparse_set.rs +++ b/benches/benches/bevy_ecs/iteration/iter_simple_foreach_sparse_set.rs @@ -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; }); } diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide.rs b/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide.rs index 4c0ee60dbec4e..7c9e7796199f8 100644 --- a/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide.rs +++ b/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide.rs @@ -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; diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide_sparse_set.rs b/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide_sparse_set.rs index f15dd5850f848..254b27a8873c5 100644 --- a/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide_sparse_set.rs +++ b/benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide_sparse_set.rs @@ -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; diff --git a/benches/benches/bevy_ecs/scheduling/running_systems.rs b/benches/benches/bevy_ecs/scheduling/running_systems.rs index 00c5b3ec0ee27..6d5d0ed3df566 100644 --- a/benches/benches/bevy_ecs/scheduling/running_systems.rs +++ b/benches/benches/bevy_ecs/scheduling/running_systems.rs @@ -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); }); } diff --git a/benches/benches/bevy_ecs/scheduling/schedule.rs b/benches/benches/bevy_ecs/scheduling/schedule.rs index 99f8b20f2bc05..a45db34c4368f 100644 --- a/benches/benches/bevy_ecs/scheduling/schedule.rs +++ b/benches/benches/bevy_ecs/scheduling/schedule.rs @@ -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); }); } diff --git a/benches/benches/bevy_ecs/world/world_get.rs b/benches/benches/bevy_ecs/world/world_get.rs index 1aba723802b0f..5daef6cd88f0d 100644 --- a/benches/benches/bevy_ecs/world/world_get.rs +++ b/benches/benches/bevy_ecs/world/world_get.rs @@ -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); diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 69ce32e8a8806..7372d92c25e0a 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -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>() - .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>() - .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] diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index f5359a25fbf07..76b9535f2a870 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -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( + &mut self, + func: &mut Func, + table: &'w Table, + rows: Range, + ) 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( + &mut self, + func: &mut Func, + archetype: &'w Archetype, + rows: Range, + ) 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( + &mut self, + mut accum: B, + func: &mut Func, + table: &'w Table, + rows: Range, + ) -> 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( + &mut self, + mut accum: B, + func: &mut Func, + archetype: &'w Archetype, + indices: Range, + ) -> 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(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, diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 2e6f4df187a40..e3564494cfcc8 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -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(); diff --git a/crates/bevy_ecs/src/query/par_iter.rs b/crates/bevy_ecs/src/query/par_iter.rs index c21931c27ebf3..2f3925593ed7f 100644 --- a/crates/bevy_ecs/src/query/par_iter.rs +++ b/crates/bevy_ecs/src/query/par_iter.rs @@ -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. diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index d4afbfe5ac9c0..c65ad65ec2706 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -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 QueryState { /// 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 QueryState { 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 QueryState { } } - /// 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 QueryState { 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 QueryState { } 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; } } diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 54f20848cbc99..bf1df191e5d70 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -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); }; } diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs index 1cc111e1374bd..f1d6048e19995 100644 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs @@ -75,16 +75,16 @@ fn main() { { let mut opt_data: Option<&Foo> = None; let mut opt_data_2: Option> = 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> = 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"); diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr index fe207ade82275..042cb246f8019 100644 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr @@ -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 diff --git a/examples/stress_tests/many_buttons.rs b/examples/stress_tests/many_buttons.rs index 3521a92f7d183..64b69902ae649 100644 --- a/examples/stress_tests/many_buttons.rs +++ b/examples/stress_tests/many_buttons.rs @@ -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()); }); }