From 019432af2e9c26dfbd6d79dda0ad107811874124 Mon Sep 17 00:00:00 2001
From: Nicola Papale <nicopap@users.noreply.github.com>
Date: Tue, 13 Jun 2023 10:47:55 +0200
Subject: [PATCH] Add get_ref to EntityRef (#8818)

# Objective

To mirror the `Ref` added as `WorldQuery`, and the `Mut` in
`EntityMut::get_mut`, we add `EntityRef::get_ref`, which retrieves `T`
with tick information, but *immutably*.

## Solution

- Add the method in question, also add it to`UnsafeEntityCell` since
this seems to be the best way of getting that information.

Also update/add safety comments to neighboring code.

---

## Changelog

- Add `EntityRef::get_ref` to get an `Option<Ref<T>>` from `EntityRef`

---------

Co-authored-by: James Liu <contact@jamessliu.com>
---
 crates/bevy_ecs/src/world/entity_ref.rs       | 12 +++++-
 .../bevy_ecs/src/world/unsafe_world_cell.rs   | 41 ++++++++++++++++---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index c97d9a2148..0a74cecfa2 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -12,7 +12,7 @@ use bevy_ptr::{OwningPtr, Ptr};
 use bevy_utils::tracing::debug;
 use std::any::TypeId;
 
-use super::unsafe_world_cell::UnsafeEntityCell;
+use super::{unsafe_world_cell::UnsafeEntityCell, Ref};
 
 /// A read-only reference to a particular [`Entity`] and all of its components
 #[derive(Copy, Clone)]
@@ -121,6 +121,16 @@ impl<'w> EntityRef<'w> {
         unsafe { self.as_unsafe_world_cell_readonly().get::<T>() }
     }
 
+    /// Gets access to the component of type `T` for the current entity,
+    /// including change detection information as a [`Ref`].
+    ///
+    /// Returns `None` if the entity does not have a component of type `T`.
+    #[inline]
+    pub fn get_ref<T: Component>(&self) -> Option<Ref<'w, T>> {
+        // SAFETY: &self implies shared access for duration of returned value
+        unsafe { self.as_unsafe_world_cell_readonly().get_ref::<T>() }
+    }
+
     /// Retrieves the change ticks for the given component. This can be useful for implementing change
     /// detection in custom runtimes.
     #[inline]
diff --git a/crates/bevy_ecs/src/world/unsafe_world_cell.rs b/crates/bevy_ecs/src/world/unsafe_world_cell.rs
index 997862981a..f95df96429 100644
--- a/crates/bevy_ecs/src/world/unsafe_world_cell.rs
+++ b/crates/bevy_ecs/src/world/unsafe_world_cell.rs
@@ -2,11 +2,11 @@
 
 #![warn(unsafe_op_in_unsafe_fn)]
 
-use super::{Mut, World, WorldId};
+use super::{Mut, Ref, World, WorldId};
 use crate::{
     archetype::{Archetype, ArchetypeComponentId, Archetypes},
     bundle::Bundles,
-    change_detection::{MutUntyped, TicksMut},
+    change_detection::{MutUntyped, Ticks, TicksMut},
     component::{
         ComponentId, ComponentStorage, ComponentTicks, Components, StorageType, Tick, TickCells,
     },
@@ -652,10 +652,10 @@ impl<'w> UnsafeEntityCell<'w> {
     #[inline]
     pub unsafe fn get<T: Component>(self) -> Option<&'w T> {
         let component_id = self.world.components().get_id(TypeId::of::<T>())?;
-
         // SAFETY:
-        // - entity location is valid
-        // - proper world access is promised by caller
+        // - `storage_type` is correct (T component_id + T::STORAGE_TYPE)
+        // - `location` is valid
+        // - proper aliasing is promised by caller
         unsafe {
             get_component(
                 self.world,
@@ -669,6 +669,36 @@ impl<'w> UnsafeEntityCell<'w> {
         }
     }
 
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeEntityCell`] has permission to access the component
+    /// - no other mutable references to the component exist at the same time
+    #[inline]
+    pub unsafe fn get_ref<T: Component>(self) -> Option<Ref<'w, T>> {
+        let last_change_tick = self.world.last_change_tick();
+        let change_tick = self.world.change_tick();
+        let component_id = self.world.components().get_id(TypeId::of::<T>())?;
+
+        // SAFETY:
+        // - `storage_type` is correct (T component_id + T::STORAGE_TYPE)
+        // - `location` is valid
+        // - proper aliasing is promised by caller
+        unsafe {
+            get_component_and_ticks(
+                self.world,
+                component_id,
+                T::Storage::STORAGE_TYPE,
+                self.entity,
+                self.location,
+            )
+            .map(|(value, cells)| Ref {
+                // SAFETY: returned component is of type T
+                value: value.deref::<T>(),
+                ticks: Ticks::from_tick_cells(cells, last_change_tick, change_tick),
+            })
+        }
+    }
+
     /// Retrieves the change ticks for the given component. This can be useful for implementing change
     /// detection in custom runtimes.
     ///
@@ -761,6 +791,7 @@ impl<'w> UnsafeEntityCell<'w> {
                 self.location,
             )
             .map(|(value, cells)| Mut {
+                // SAFETY: returned component is of type T
                 value: value.assert_unique().deref_mut::<T>(),
                 ticks: TicksMut::from_tick_cells(cells, last_change_tick, change_tick),
             })