[Merged by Bors] - Speed up Query::get_many and add benchmarks

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

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<const N: usize>(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::<Query<&Table>>::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::<Query<&Sparse>>::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);
+            });
+        });
+    }
+}

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<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         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<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
             }
         }
 
-        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`].