diff --git a/benches/benches/bevy_ecs/world/mod.rs b/benches/benches/bevy_ecs/world/mod.rs index 8505087e48..f594d082d7 100644 --- a/benches/benches/bevy_ecs/world/mod.rs +++ b/benches/benches/bevy_ecs/world/mod.rs @@ -27,4 +27,7 @@ criterion_group!( query_get_component_simple, query_get_component, query_get, + query_get_many::<2>, + query_get_many::<5>, + query_get_many::<10>, ); diff --git a/benches/benches/bevy_ecs/world/world_get.rs b/benches/benches/bevy_ecs/world/world_get.rs index 2a4e7337fb..536b5a7b8b 100644 --- a/benches/benches/bevy_ecs/world/world_get.rs +++ b/benches/benches/bevy_ecs/world/world_get.rs @@ -402,3 +402,58 @@ pub fn query_get(criterion: &mut Criterion) { group.finish(); } + +pub fn query_get_many(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group(&format!("query_get_many_{N}")); + group.warm_up_time(std::time::Duration::from_millis(500)); + group.measurement_time(std::time::Duration::from_secs(2 * N as u64)); + + for entity_count in RANGE.map(|i| i * 10_000) { + group.bench_function(format!("{}_calls_table", entity_count), |bencher| { + let mut world = World::default(); + let mut entity_groups: Vec<_> = (0..entity_count) + .map(|_| [(); N].map(|_| world.spawn(Table::default()).id())) + .collect(); + entity_groups.shuffle(&mut deterministic_rand()); + + let mut query = SystemState::>::new(&mut world); + let query = query.get(&world); + + bencher.iter(|| { + let mut count = 0; + for comp in entity_groups + .iter() + .filter_map(|&ids| query.get_many(ids).ok()) + { + black_box(comp); + count += 1; + black_box(count); + } + assert_eq!(black_box(count), entity_count); + }); + }); + group.bench_function(format!("{}_calls_sparse", entity_count), |bencher| { + let mut world = World::default(); + let mut entity_groups: Vec<_> = (0..entity_count) + .map(|_| [(); N].map(|_| world.spawn(Sparse::default()).id())) + .collect(); + entity_groups.shuffle(&mut deterministic_rand()); + + let mut query = SystemState::>::new(&mut world); + let query = query.get(&world); + + bencher.iter(|| { + let mut count = 0; + for comp in entity_groups + .iter() + .filter_map(|&ids| query.get_many(ids).ok()) + { + black_box(comp); + count += 1; + black_box(count); + } + assert_eq!(black_box(count), entity_count); + }); + }); + } +} diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 3a22ccb328..0a26c16eba 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -11,7 +11,7 @@ use bevy_tasks::ComputeTaskPool; #[cfg(feature = "trace")] use bevy_utils::tracing::Instrument; use fixedbitset::FixedBitSet; -use std::{borrow::Borrow, fmt}; +use std::{borrow::Borrow, fmt, mem::MaybeUninit}; use super::{NopWorldQuery, QueryItem, QueryManyIter, ROQueryItem, ReadOnlyWorldQuery}; @@ -438,24 +438,22 @@ impl QueryState { last_change_tick: u32, change_tick: u32, ) -> Result<[ROQueryItem<'w, Q>; N], QueryEntityError> { - // SAFETY: fetch is read-only - // and world must be validated - let array_of_results = entities.map(|entity| { - self.as_readonly() - .get_unchecked_manual(world, entity, last_change_tick, change_tick) - }); + let mut values = [(); N].map(|_| MaybeUninit::uninit()); - // TODO: Replace with TryMap once https://github.com/rust-lang/rust/issues/79711 is stabilized - // If any of the get calls failed, bubble up the error - for result in &array_of_results { - match result { - Ok(_) => (), - Err(error) => return Err(*error), - } + for (value, entity) in std::iter::zip(&mut values, entities) { + // SAFETY: fetch is read-only + // and world must be validated + let item = self.as_readonly().get_unchecked_manual( + world, + entity, + last_change_tick, + change_tick, + )?; + *value = MaybeUninit::new(item); } - // Since we have verified that all entities are present, we can safely unwrap - Ok(array_of_results.map(|result| result.unwrap())) + // SAFETY: Each value has been fully initialized. + Ok(values.map(|x| x.assume_init())) } /// Gets the query results for the given [`World`] and array of [`Entity`], where the last change and @@ -484,19 +482,15 @@ impl QueryState { } } - let array_of_results = entities - .map(|entity| self.get_unchecked_manual(world, entity, last_change_tick, change_tick)); + let mut values = [(); N].map(|_| MaybeUninit::uninit()); - // If any of the get calls failed, bubble up the error - for result in &array_of_results { - match result { - Ok(_) => (), - Err(error) => return Err(*error), - } + for (value, entity) in std::iter::zip(&mut values, entities) { + let item = self.get_unchecked_manual(world, entity, last_change_tick, change_tick)?; + *value = MaybeUninit::new(item); } - // Since we have verified that all entities are present, we can safely unwrap - Ok(array_of_results.map(|result| result.unwrap())) + // SAFETY: Each value has been fully initialized. + Ok(values.map(|x| x.assume_init())) } /// Returns an [`Iterator`] over the query results for the given [`World`].