From 7166a28bafa669eac40f5864cce8a150b330dd5f Mon Sep 17 00:00:00 2001
From: Will Crichton <wcrichto@cs.stanford.edu>
Date: Fri, 22 Jan 2021 18:15:08 -0500
Subject: [PATCH] Enable dynamic mutable access to component data  (#1284)

* Enable dynamic mutable access to component data

* Add clippy allowance, more documentation
---
 crates/bevy_ecs/src/core/archetype.rs     |  7 +------
 crates/bevy_reflect/src/impls/bevy_ecs.rs | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_ecs/src/core/archetype.rs b/crates/bevy_ecs/src/core/archetype.rs
index cc10e9d42a..21d44ee5ac 100644
--- a/crates/bevy_ecs/src/core/archetype.rs
+++ b/crates/bevy_ecs/src/core/archetype.rs
@@ -223,12 +223,7 @@ impl Archetype {
 
     /// # Safety
     /// `index` must be in-bounds
-    pub(crate) unsafe fn get_dynamic(
-        &self,
-        ty: TypeId,
-        size: usize,
-        index: usize,
-    ) -> Option<NonNull<u8>> {
+    pub unsafe fn get_dynamic(&self, ty: TypeId, size: usize, index: usize) -> Option<NonNull<u8>> {
         debug_assert!(index < self.len);
         Some(NonNull::new_unchecked(
             (*self.data.get())
diff --git a/crates/bevy_reflect/src/impls/bevy_ecs.rs b/crates/bevy_reflect/src/impls/bevy_ecs.rs
index 4be4a77bac..a50e61f1a3 100644
--- a/crates/bevy_reflect/src/impls/bevy_ecs.rs
+++ b/crates/bevy_reflect/src/impls/bevy_ecs.rs
@@ -10,6 +10,7 @@ pub struct ReflectComponent {
     add_component: fn(&mut World, resources: &Resources, Entity, &dyn Reflect),
     apply_component: fn(&mut World, Entity, &dyn Reflect),
     reflect_component: unsafe fn(&Archetype, usize) -> &dyn Reflect,
+    reflect_component_mut: unsafe fn(&Archetype, usize) -> &mut dyn Reflect,
     copy_component: fn(&World, &mut World, &Resources, Entity, Entity),
 }
 
@@ -38,6 +39,20 @@ impl ReflectComponent {
         (self.reflect_component)(archetype, entity_index)
     }
 
+    /// # Safety
+    /// This does not do bound checks on entity_index. You must make sure entity_index is within bounds before calling.
+    /// 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 a thread-local system to avoid sharing across threads.
+    /// * Don't call this method more than once in the same scope for a given component.
+    #[allow(clippy::mut_from_ref)]
+    pub unsafe fn reflect_component_mut<'a>(
+        &self,
+        archetype: &'a Archetype,
+        entity_index: usize,
+    ) -> &'a mut dyn Reflect {
+        (self.reflect_component_mut)(archetype, entity_index)
+    }
+
     pub fn copy_component(
         &self,
         source_world: &World,
@@ -87,6 +102,13 @@ impl<C: Component + Reflect + FromResources> FromType<C> for ReflectComponent {
                     ptr.as_ref().unwrap()
                 }
             },
+            reflect_component_mut: |archetype, index| {
+                unsafe {
+                    // the type has been looked up by the caller, so this is safe
+                    let ptr = archetype.get::<C>().unwrap().as_ptr().add(index);
+                    &mut *ptr
+                }
+            },
         }
     }
 }