From e0a6c718014d43cefcf9cd33e228c018ab6f48ea Mon Sep 17 00:00:00 2001
From: Jakob Hellermann <jakob.hellermann@protonmail.com>
Date: Fri, 9 Sep 2022 11:03:21 +0200
Subject: [PATCH] add EntityRef::get_unchecked_mut_by_id

---
 crates/bevy_ecs/src/world/entity_ref.rs | 35 ++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index e4b7e072d4ff8..54c603e38a65a 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -126,6 +126,31 @@ impl<'w> EntityRef<'w> {
         // SAFETY: entity_location is valid, component_id is valid as checked by the line above
         unsafe { get_component(self.world, component_id, self.entity, self.location) }
     }
+
+    /// Gets the component of the given [`ComponentId`] from the entity.
+    ///
+    /// **You should prefer to use the typed API where possible and only
+    /// use this in cases where the actual component types are not known at
+    /// compile time.**
+    ///
+    /// Unlike [`EntityRef::get`], this returns a raw pointer to the component,
+    /// which is only valid while the `'w` borrow of the lifetime is active.
+    ///
+    /// # Safety
+    ///
+    /// - The returned reference must never alias a mutable borrow of this component.
+    /// - The returned reference must not be used after this component is moved which
+    ///   may happen from **any** `insert_component`, `remove_component` or `despawn`
+    ///   operation on this world (non-exhaustive list).
+    #[inline]
+    pub unsafe fn get_unchecked_mut_by_id(
+        &self,
+        component_id: ComponentId,
+    ) -> Option<MutUntyped<'w>> {
+        self.world.components().get_info(component_id)?;
+        // SAFETY: entity_location is valid, component_id is valid as checked by the line above, world access is promised by the caller
+        unsafe { get_mut_by_id(self.world, self.entity, self.location, component_id) }
+    }
 }
 
 impl<'w> From<EntityMut<'w>> for EntityRef<'w> {
@@ -555,7 +580,7 @@ impl<'w> EntityMut<'w> {
     #[inline]
     pub fn get_mut_by_id(&mut self, component_id: ComponentId) -> Option<MutUntyped<'_>> {
         self.world.components().get_info(component_id)?;
-        // SAFETY: entity_location is valid, component_id is valid as checked by the line above
+        // SAFETY: entity_location is valid, component_id is valid as checked by the line above, world access is unique
         unsafe { get_mut_by_id(self.world, self.entity, self.location, component_id) }
     }
 }
@@ -889,15 +914,17 @@ pub(crate) unsafe fn get_mut<T: Component>(
     )
 }
 
-// SAFETY: EntityLocation must be valid, component_id must be valid
+// SAFETY:
+// - EntityLocation must be valid, component_id must be valid
+// - world access to the component must be valid, either because the caller has a `&mut` world or it synchronizes access like systems do
 #[inline]
 pub(crate) unsafe fn get_mut_by_id(
-    world: &mut World,
+    world: &World,
     entity: Entity,
     location: EntityLocation,
     component_id: ComponentId,
 ) -> Option<MutUntyped> {
-    // SAFETY: world access is unique, entity location and component_id required to be valid
+    // SAFETY: world access promised by the caller, entity location and component_id required to be valid
     get_component_and_ticks(world, component_id, entity, location).map(|(value, ticks)| {
         MutUntyped {
             value: value.assert_unique(),