Optimize World::get_entity_mut for large entity slices

benches/benches/bevy_ecs/world/entity_fetch.rs

@@ -0,0 +1,23 @@
+use bevy_ecs::{entity::Entity, world::World};
+use core::hint::black_box;
+use criterion::{BenchmarkId, Criterion};
+use std::time::Duration;
+
+pub fn get_entity_mut_slice(criterion: &mut Criterion) {
+    let mut group = criterion.benchmark_group("get_entity_mut_slice");
+    group.warm_up_time(Duration::from_millis(500));
+    group.measurement_time(Duration::from_secs(4));
+
+    for size in [10, 20, 40, 60, 80, 100, 200, 400, 600, 800, 1000, 2000].iter() {
+        group.bench_with_input(BenchmarkId::new("size", size), size, |b, &size| {
+            let mut world = World::new();
+            let entities: Vec<Entity> = (0..size).map(|_| world.spawn_empty().id()).collect();
+
+            b.iter(|| {
+                // This triggers `WorldEntityFetch for &'_ [Entity]`
+                let _ = world.get_entity_mut(black_box(entities.as_slice()));
+            });
+        });
+    }
+    group.finish();
+}

benches/benches/bevy_ecs/world/mod.rs

@@ -2,6 +2,7 @@ mod commands;
 mod despawn;
 mod despawn_recursive;
 mod entity_allocator;
+mod entity_fetch;
 mod entity_hash;
 mod spawn;
 mod world_get;
@@ -11,6 +12,7 @@ use criterion::criterion_group;
 use despawn::*;
 use despawn_recursive::*;
 use entity_allocator::*;
+use entity_fetch::*;
 use entity_hash::*;
 use spawn::*;
 use world_get::*;
@@ -44,4 +46,5 @@ criterion_group!(
     query_get_components_mut_32,
     entity_set_build_and_lookup,
     entity_allocator_benches,
+    get_entity_mut_slice
 );

crates/bevy_ecs/src/world/entity_fetch.rs

@@ -302,14 +302,7 @@ unsafe impl<const N: usize> WorldEntityFetch for &'_ [Entity; N] {
         self,
         cell: UnsafeWorldCell<'_>,
     ) -> Result<Self::Mut<'_>, EntityMutableFetchError> {
-        // Check for duplicate entities.
-        for i in 0..self.len() {
-            for j in 0..i {
-                if self[i] == self[j] {
-                    return Err(EntityMutableFetchError::AliasedMutability(self[i]));
-                }
-            }
-        }
+        check_duplicate_entities(self)?;
 
         let mut refs = [const { MaybeUninit::uninit() }; N];
         for (r, &id) in core::iter::zip(&mut refs, self) {
@@ -364,14 +357,7 @@ unsafe impl WorldEntityFetch for &'_ [Entity] {
         self,
         cell: UnsafeWorldCell<'_>,
     ) -> Result<Self::Mut<'_>, EntityMutableFetchError> {
-        // Check for duplicate entities.
-        for i in 0..self.len() {
-            for j in 0..i {
-                if self[i] == self[j] {
-                    return Err(EntityMutableFetchError::AliasedMutability(self[i]));
-                }
-            }
-        }
+        check_duplicate_entities(self)?;
 
         let mut refs = Vec::with_capacity(self.len());
         for &id in self {
@@ -441,3 +427,28 @@ unsafe impl WorldEntityFetch for &'_ EntityHashSet {
         unsafe { self.fetch_mut(cell) }
     }
 }
+
+/// Checks for duplicate entities in a slice.
+#[inline]
+fn check_duplicate_entities(entities: &[Entity]) -> Result<(), EntityMutableFetchError> {
+    const THRESHOLD: usize = 40;
+
+    if entities.len() <= THRESHOLD {
+        for i in 0..entities.len() {
+            for j in 0..i {
+                if entities[i] == entities[j] {
+                    return Err(EntityMutableFetchError::AliasedMutability(entities[i]));
+                }
+            }
+        }
+    } else {
+        let mut seen = EntityHashSet::with_capacity(entities.len());
+        for &id in entities {
+            if !seen.insert(id) {
+                return Err(EntityMutableFetchError::AliasedMutability(id));
+            }
+        }
+    }
+
+    Ok(())
+}