use InteriorMutableWorld for

ReflectComponent/ReflectResource/ReflectAsset

Author: jakobhellermann

crates/bevy_asset/src/reflect.rs

@@ -1,6 +1,6 @@
 use std::any::{Any, TypeId};
 
-use bevy_ecs::world::World;
+use bevy_ecs::world::{interior_mutable_world::InteriorMutableWorld, World};
 use bevy_reflect::{FromReflect, FromType, Reflect, Uuid};
 
 use crate::{Asset, Assets, Handle, HandleId, HandleUntyped};
@@ -18,8 +18,11 @@ pub struct ReflectAsset {
     assets_resource_type_id: TypeId,
 
     get: fn(&World, HandleUntyped) -> Option<&dyn Reflect>,
-    get_mut: fn(&mut World, HandleUntyped) -> Option<&mut dyn Reflect>,
-    get_unchecked_mut: unsafe fn(&World, HandleUntyped) -> Option<&mut dyn Reflect>,
+    // SAFETY:
+    // - may only be called with a [`IteriorMutableWorld`] which can be used to access the corresponding `Assets<T>` resource mutably
+    // - may only be used to access **at most one** access at once
+    get_unchecked_mut:
+        unsafe fn(InteriorMutableWorld<'_>, HandleUntyped) -> Option<&mut dyn Reflect>,
     add: fn(&mut World, &dyn Reflect) -> HandleUntyped,
     set: fn(&mut World, HandleUntyped, &dyn Reflect) -> HandleUntyped,
     len: fn(&World) -> usize,
@@ -54,10 +57,11 @@ impl ReflectAsset {
         world: &'w mut World,
         handle: HandleUntyped,
     ) -> Option<&'w mut dyn Reflect> {
-        (self.get_mut)(world, handle)
+        // SAFETY: unique world access
+        unsafe { (self.get_unchecked_mut)(world.as_interior_mutable(), handle) }
     }
 
-    /// Equivalent of [`Assets::get_mut`], but does not require a mutable reference to the world.
+    /// Equivalent of [`Assets::get_mut`], but works with a [`InteriorMutableWorld`].
     ///
     /// Only use this method when you have ensured that you are the *only* one with access to the [`Assets`] resource of the asset type.
     /// Furthermore, this does *not* allow you to have look up two distinct handles,
@@ -81,12 +85,11 @@ impl ReflectAsset {
     /// # 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 when you have exclusive access to the world
-    /// (or use a scheduler that enforces unique access to the `Assets` resource).
+    /// * Only call this method if you know that the [`InteriorMutableWorld`] may be used to access the corresponding `Assets<T>`
     /// * Don't call this method more than once in the same scope.
     pub unsafe fn get_unchecked_mut<'w>(
         &self,
-        world: &'w World,
+        world: InteriorMutableWorld<'w>,
         handle: HandleUntyped,
     ) -> Option<&'w mut dyn Reflect> {
         // SAFETY: requirements are deferred to the caller
@@ -140,19 +143,8 @@ impl<A: Asset + FromReflect> FromType<A> for ReflectAsset {
                 let asset = assets.get(&handle.typed());
                 asset.map(|asset| asset as &dyn Reflect)
             },
-            get_mut: |world, handle| {
-                let assets = world.resource_mut::<Assets<A>>().into_inner();
-                let asset = assets.get_mut(&handle.typed());
-                asset.map(|asset| asset as &mut dyn Reflect)
-            },
             get_unchecked_mut: |world, handle| {
-                let assets = unsafe {
-                    world
-                        .as_interior_mutable()
-                        .get_resource_mut::<Assets<A>>()
-                        .unwrap()
-                        .into_inner()
-                };
+                let assets = unsafe { world.get_resource_mut::<Assets<A>>().unwrap().into_inner() };
                 let asset = assets.get_mut(&handle.typed());
                 asset.map(|asset| asset as &mut dyn Reflect)
             },

crates/bevy_ecs/src/reflect.rs

@@ -5,7 +5,7 @@ use crate::{
     component::Component,
     entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
     system::Resource,
-    world::{FromWorld, World},
+    world::{interior_mutable_world::InteriorMutableWorld, FromWorld, World},
 };
 use bevy_reflect::{
     impl_from_reflect_value, impl_reflect_value, FromType, Reflect, ReflectDeserialize,
@@ -52,7 +52,8 @@ pub struct ReflectComponentFns {
     /// Function pointer implementing [`ReflectComponent::reflect()`].
     pub reflect: fn(&World, Entity) -> Option<&dyn Reflect>,
     /// Function pointer implementing [`ReflectComponent::reflect_mut()`].
-    pub reflect_mut: unsafe fn(&World, Entity) -> Option<Mut<dyn Reflect>>,
+    /// The function may only be called with a InteriorMutableWorld that can be used to access the relevant component on the given entity
+    pub reflect_mut: unsafe fn(InteriorMutableWorld<'_>, Entity) -> Option<Mut<'_, dyn Reflect>>,
     /// Function pointer implementing [`ReflectComponent::copy()`].
     pub copy: fn(&World, &mut World, Entity, Entity),
 }
@@ -117,20 +118,20 @@ impl ReflectComponent {
         entity: Entity,
     ) -> Option<Mut<'a, dyn Reflect>> {
         // SAFETY: unique world access
-        unsafe { (self.0.reflect_mut)(world, entity) }
+        unsafe { (self.0.reflect_mut)(world.as_interior_mutable(), 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).
+    /// * Only call this method with a [`InteriorMutableWorld`] that may be used to access the component on the entity `entity`
     /// * 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,
+        world: InteriorMutableWorld<'a>,
         entity: Entity,
     ) -> Option<Mut<'a, dyn Reflect>> {
+        // SAFETY: safety requirements deferred to caller
         (self.0.reflect_mut)(world, entity)
     }
 
@@ -212,6 +213,9 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
                 // 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 {
+                    // SAFETY: entity access through the InteriorMutableWorld is not implemented yet.
+                    // The following code only accesses the component `C` through the entity `entity`
+                    let world = world.world();
                     world
                         .get_entity(entity)?
                         .get_unchecked_mut::<C>(world.last_change_tick(), world.read_change_tick())
@@ -265,7 +269,7 @@ pub struct ReflectResourceFns {
     /// Function pointer implementing [`ReflectResource::reflect()`].
     pub reflect: fn(&World) -> Option<&dyn Reflect>,
     /// Function pointer implementing [`ReflectResource::reflect_unchecked_mut()`].
-    pub reflect_unchecked_mut: unsafe fn(&World) -> Option<Mut<dyn Reflect>>,
+    pub reflect_unchecked_mut: unsafe fn(InteriorMutableWorld<'_>) -> Option<Mut<'_, dyn Reflect>>,
     /// Function pointer implementing [`ReflectResource::copy()`].
     pub copy: fn(&World, &mut World),
 }
@@ -314,19 +318,18 @@ impl ReflectResource {
     /// Gets the value of this [`Resource`] type from the world as a mutable reflected reference.
     pub fn reflect_mut<'a>(&self, world: &'a mut World) -> Option<Mut<'a, dyn Reflect>> {
         // SAFETY: unique world access
-        unsafe { (self.0.reflect_unchecked_mut)(world) }
+        unsafe { (self.0.reflect_unchecked_mut)(world.as_interior_mutable()) }
     }
 
     /// # 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).
+    /// * Only call this method with a [`InteriorMutableWorld`] which can be used to access the resource.
     /// * Don't call this method more than once in the same scope for a given [`Resource`].
-    pub unsafe fn reflect_unchecked_mut<'a>(
+    pub unsafe fn reflect_unchecked_mut<'w>(
         &self,
-        world: &'a World,
-    ) -> Option<Mut<'a, dyn Reflect>> {
+        world: InteriorMutableWorld<'w>,
+    ) -> Option<Mut<'w, dyn Reflect>> {
         // SAFETY: caller promises to uphold uniqueness guarantees
         (self.0.reflect_unchecked_mut)(world)
     }
@@ -385,13 +388,10 @@ impl<C: Resource + Reflect + FromWorld> FromType<C> for ReflectResource {
                 // SAFETY: all usages of `reflect_unchecked_mut` guarantee that there is either a single mutable
                 // reference or multiple immutable ones alive at any given point
                 unsafe {
-                    world
-                        .as_interior_mutable()
-                        .get_resource_mut::<C>()
-                        .map(|res| Mut {
-                            value: res.value as &mut dyn Reflect,
-                            ticks: res.ticks,
-                        })
+                    world.get_resource_mut::<C>().map(|res| Mut {
+                        value: res.value as &mut dyn Reflect,
+                        ticks: res.ticks,
+                    })
                 }
             },
             copy: |source_world, destination_world| {