From b83c0e106e6905d8b3a6c744bfa48352f0e8d912 Mon Sep 17 00:00:00 2001
From: Sou1gh0st <xiuyutong1994@163.com>
Date: Tue, 12 Nov 2024 02:46:47 +0800
Subject: [PATCH] Add EntityMut::get_mut_by_id_unchecked (#16210)

# Objective

- Fixes: #15603

## Solution

- Add an unsafe `get_mut_by_id_unchecked` to `EntityMut` that borrows
&self instead of &mut self, thereby allowing access to multiple
components simultaneously.

## Testing

- a unit test function `get_mut_by_id_unchecked` was added.

---------

Co-authored-by: Mike <mike.hsu@gmail.com>
---
 crates/bevy_ecs/src/world/entity_ref.rs | 51 +++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index 839833dc02..8917b924b5 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -730,6 +730,34 @@ impl<'w> EntityMut<'w> {
         unsafe { component_ids.fetch_mut(self.0) }
     }
 
+    /// Returns [untyped mutable reference](MutUntyped) to component for
+    /// the current entity, based on the given [`ComponentId`].
+    ///
+    /// Unlike [`EntityMut::get_mut_by_id`], this method borrows &self instead of
+    /// &mut self, allowing the caller to access multiple components simultaneously.
+    ///
+    /// # Errors
+    ///
+    /// - Returns [`EntityComponentError::MissingComponent`] if the entity does
+    ///   not have a component.
+    /// - Returns [`EntityComponentError::AliasedMutability`] if a component
+    ///   is requested multiple times.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeEntityCell`] has permission to access the component mutably
+    /// - no other references to the component exist at the same time
+    #[inline]
+    pub unsafe fn get_mut_by_id_unchecked<F: DynamicComponentFetch>(
+        &self,
+        component_ids: F,
+    ) -> Result<F::Mut<'_>, EntityComponentError> {
+        // SAFETY:
+        // - The caller must ensure simultaneous access is limited
+        // - to components that are mutually independent.
+        unsafe { component_ids.fetch_mut(self.0) }
+    }
+
     /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped)
     /// to component(s) with lifetime `'w` for the current entity, based on the
     /// given [`ComponentId`]s.
@@ -4425,4 +4453,27 @@ mod tests {
                 .map(|_| { unreachable!() })
         );
     }
+
+    #[test]
+    fn get_mut_by_id_unchecked() {
+        let mut world = World::default();
+        let e1 = world.spawn((X(7), Y(10))).id();
+        let x_id = world.register_component::<X>();
+        let y_id = world.register_component::<Y>();
+
+        let e1_mut = &world.get_entity_mut([e1]).unwrap()[0];
+        // SAFETY: The entity e1 contains component X.
+        let x_ptr = unsafe { e1_mut.get_mut_by_id_unchecked(x_id) }.unwrap();
+        // SAFETY: The entity e1 contains component Y, with components X and Y being mutually independent.
+        let y_ptr = unsafe { e1_mut.get_mut_by_id_unchecked(y_id) }.unwrap();
+
+        // SAFETY: components match the id they were fetched with
+        let x_component = unsafe { x_ptr.into_inner().deref_mut::<X>() };
+        x_component.0 += 1;
+        // SAFETY: components match the id they were fetched with
+        let y_component = unsafe { y_ptr.into_inner().deref_mut::<Y>() };
+        y_component.0 -= 1;
+
+        assert_eq!((&mut X(8), &mut Y(9)), (x_component, y_component));
+    }
 }