From eb129b7b2d3d1d95c838456c0a494bbfb66d12f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20=C5=81abor?= <jacob.labor@gmail.com>
Date: Thu, 19 Jan 2023 18:31:11 +0100
Subject: [PATCH] Review items

---
 crates/bevy_ecs/src/reflect.rs                | 39 ++++++++++++++++---
 .../bevy_scene/src/dynamic_scene_builder.rs   |  2 +-
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_ecs/src/reflect.rs b/crates/bevy_ecs/src/reflect.rs
index d1568c0e605a44..27978563354511 100644
--- a/crates/bevy_ecs/src/reflect.rs
+++ b/crates/bevy_ecs/src/reflect.rs
@@ -50,11 +50,13 @@ pub struct ReflectComponentFns {
     /// Function pointer implementing [`ReflectComponent::remove()`].
     pub remove: fn(&mut EntityMut),
     /// Function pointer implementing [`ReflectComponent::contains()`].
-    pub contains: fn(&EntityRef) -> bool,
+    pub contains: fn(EntityRef) -> bool,
     /// Function pointer implementing [`ReflectComponent::reflect()`].
-    pub reflect: for<'a> fn(&'a EntityRef) -> Option<&'a dyn Reflect>,
+    pub reflect: fn(EntityRef) -> Option<&dyn Reflect>,
     /// Function pointer implementing [`ReflectComponent::reflect_mut()`].
-    pub reflect_mut: for<'a> fn(&'a mut EntityMut) -> Option<Mut<'a, dyn Reflect>>,
+    pub reflect_mut: for<'a> fn(&'a mut EntityMut<'a>) -> Option<Mut<'a, dyn Reflect>>,
+    /// Function pointer implementing [`ReflectComponent::reflect_unchecked_mut()`].
+    pub reflect_unchecked_mut: fn(&World, Entity) -> Option<Mut<dyn Reflect>>,
     /// Function pointer implementing [`ReflectComponent::copy()`].
     pub copy: fn(&World, &mut World, Entity, Entity),
 }
@@ -108,12 +110,12 @@ impl ReflectComponent {
     }
 
     /// Returns whether entity contains this [`Component`]
-    pub fn contains(&self, entity: &EntityRef) -> bool {
+    pub fn contains(&self, entity: EntityRef) -> bool {
         (self.0.contains)(entity)
     }
 
     /// Gets the value of this [`Component`] type from the entity as a reflected reference.
-    pub fn reflect<'a>(&self, entity: &'a EntityRef<'a>) -> Option<&'a dyn Reflect> {
+    pub fn reflect<'a>(&self, entity: EntityRef<'a>) -> Option<&'a dyn Reflect> {
         (self.0.reflect)(entity)
     }
 
@@ -122,6 +124,20 @@ impl ReflectComponent {
         (self.0.reflect_mut)(entity)
     }
 
+    /// # Safety
+    /// This method does not prevent you from having two mutable pointers to the same data,
+    /// violating Rust's aliasing rules. To avoid this:
+    /// * Only call this method in an exclusive system to avoid sharing across threads (or use a
+    ///   scheduler that enforces safe memory access).
+    /// * Don't call this method more than once in the same scope for a given [`Component`].
+    pub unsafe fn reflect_unchecked_mut<'a>(
+        &self,
+        world: &'a World,
+        entity: Entity,
+    ) -> Option<Mut<'a, dyn Reflect>> {
+        (self.0.reflect_unchecked_mut)(world, entity)
+    }
+
     /// Gets the value of this [`Component`] type from entity from `source_world` and [applies](Self::apply()) it to the value of this [`Component`] type in entity in `destination_world`.
     ///
     /// # Panics
@@ -198,6 +214,19 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
                     ticks: c.ticks,
                 })
             },
+            reflect_unchecked_mut: |world, entity| {
+                // SAFETY: reflect_mut is an unsafe function pointer used by `reflect_unchecked_mut` which promises to never
+                // produce aliasing mutable references, and reflect_mut, which has mutable world access
+                unsafe {
+                    world
+                        .get_entity(entity)?
+                        .get_unchecked_mut::<C>(world.last_change_tick(), world.read_change_tick())
+                        .map(|c| Mut {
+                            value: c.value as &mut dyn Reflect,
+                            ticks: c.ticks,
+                        })
+                }
+            },
         })
     }
 }
diff --git a/crates/bevy_scene/src/dynamic_scene_builder.rs b/crates/bevy_scene/src/dynamic_scene_builder.rs
index e688887564c04d..1187f087b68096 100644
--- a/crates/bevy_scene/src/dynamic_scene_builder.rs
+++ b/crates/bevy_scene/src/dynamic_scene_builder.rs
@@ -132,7 +132,7 @@ impl<'w> DynamicSceneBuilder<'w> {
                     .get_info(component_id)
                     .and_then(|info| type_registry.get(info.type_id().unwrap()))
                     .and_then(|registration| registration.data::<ReflectComponent>())
-                    .and_then(|reflect_component| reflect_component.reflect(&entity));
+                    .and_then(|reflect_component| reflect_component.reflect(entity));
 
                 if let Some(reflect_component) = reflect_component {
                     entry.components.push(reflect_component.clone_value());