[Merged by Bors] - add UnsafeWorldCell abstraction

crates/bevy_asset/src/reflect.rs

@@ -1,6 +1,6 @@
 use std::any::{Any, TypeId};
 
-use bevy_ecs::world::World;
+use bevy_ecs::world::{unsafe_world_cell::UnsafeWorldCell, World};
 use bevy_reflect::{FromReflect, FromType, Reflect, Uuid};
 
 use crate::{Asset, Assets, Handle, HandleId, HandleUntyped};
@@ -18,8 +18,10 @@ 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(UnsafeWorldCell<'_>, 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 +56,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_unsafe_world_cell(), handle) }
     }
 
-    /// Equivalent of [`Assets::get_mut`], but does not require a mutable reference to the world.
+    /// Equivalent of [`Assets::get_mut`], but works with an [`UnsafeWorldCell`].
     ///
     /// 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,
@@ -67,11 +70,12 @@ impl ReflectAsset {
     /// # use bevy_asset::{ReflectAsset, HandleUntyped};
     /// # use bevy_ecs::prelude::World;
     /// # let reflect_asset: ReflectAsset = unimplemented!();
-    /// # let world: World = unimplemented!();
+    /// # let mut world: World = unimplemented!();
     /// # let handle_1: HandleUntyped = unimplemented!();
     /// # let handle_2: HandleUntyped = unimplemented!();
-    /// let a = unsafe { reflect_asset.get_unchecked_mut(&world, handle_1).unwrap() };
-    /// let b = unsafe { reflect_asset.get_unchecked_mut(&world, handle_2).unwrap() };
+    /// let unsafe_world_cell = world.as_unsafe_world_cell();
+    /// let a = unsafe { reflect_asset.get_unchecked_mut(unsafe_world_cell, handle_1).unwrap() };
+    /// let b = unsafe { reflect_asset.get_unchecked_mut(unsafe_world_cell, handle_2).unwrap() };
     /// // ^ not allowed, two mutable references through the same asset resource, even though the
     /// // handles are distinct
     ///
@@ -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 [`UnsafeWorldCell`] 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: UnsafeWorldCell<'w>,
         handle: HandleUntyped,
     ) -> Option<&'w mut dyn Reflect> {
         // SAFETY: requirements are deferred to the caller
@@ -140,18 +143,10 @@ 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
-                        .get_resource_unchecked_mut::<Assets<A>>()
-                        .unwrap()
-                        .into_inner()
-                };
+                // SAFETY: `get_unchecked_mut` must be callied with `UnsafeWorldCell` having access to `Assets<A>`,
+                // and must ensure to only have at most one reference to it live at all times.
+                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::{unsafe_world_cell::UnsafeWorldCell, FromWorld, World},
 };
 use bevy_reflect::{
     impl_from_reflect_value, impl_reflect_value, FromType, Reflect, ReflectDeserialize,
@@ -52,7 +52,10 @@ 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>>,
+    ///
+    /// # Safety
+    /// The function may only be called with an [`UnsafeWorldCell`] that can be used to mutably access the relevant component on the given entity.
+    pub reflect_mut: unsafe fn(UnsafeWorldCell<'_>, Entity) -> Option<Mut<'_, dyn Reflect>>,
     /// Function pointer implementing [`ReflectComponent::copy()`].
     pub copy: fn(&World, &mut World, Entity, Entity),
 }
@@ -117,20 +120,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_unsafe_world_cell(), 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 [`UnsafeWorldCell`] that may be used to mutably 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: UnsafeWorldCell<'a>,
         entity: Entity,
     ) -> Option<Mut<'a, dyn Reflect>> {
+        // SAFETY: safety requirements deferred to caller
         (self.0.reflect_mut)(world, entity)
     }
 
@@ -209,16 +212,14 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
                     .map(|c| c as &dyn Reflect)
             },
             reflect_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
+                // SAFETY: reflect_mut is an unsafe function pointer used by
+                // 1. `reflect_unchecked_mut` which must be called with an UnsafeWorldCell with access to the the component `C` on the `entity`, and
+                // 2. `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,
-                        })
+                    world.get_entity(entity)?.get_mut::<C>().map(|c| Mut {
+                        value: c.value as &mut dyn Reflect,
+                        ticks: c.ticks,
+                    })
                 }
             },
         })
@@ -265,7 +266,10 @@ 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>>,
+    ///
+    /// # Safety
+    /// The function may only be called with an [`UnsafeWorldCell`] that can be used to mutably access the relevant resource.
+    pub reflect_unchecked_mut: unsafe fn(UnsafeWorldCell<'_>) -> 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_unsafe_world_cell()) }
     }
 
     /// # 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 an [`UnsafeWorldCell`] which can be used to mutably 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: UnsafeWorldCell<'w>,
+    ) -> Option<Mut<'w, dyn Reflect>> {
         // SAFETY: caller promises to uphold uniqueness guarantees
         (self.0.reflect_unchecked_mut)(world)
     }
@@ -385,7 +388,7 @@ 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.get_resource_unchecked_mut::<C>().map(|res| Mut {
+                    world.get_resource_mut::<C>().map(|res| Mut {
                         value: res.value as &mut dyn Reflect,
                         ticks: res.ticks,
                     })

crates/bevy_ecs/src/system/query.rs

@@ -1134,6 +1134,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
         }
         let world = self.world;
         let entity_ref = world
+            .as_unsafe_world_cell_migration_internal()
             .get_entity(entity)
             .ok_or(QueryComponentError::NoSuchEntity)?;
         let component_id = world
@@ -1150,7 +1151,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
             .has_write(archetype_component)
         {
             entity_ref
-                .get_unchecked_mut::<T>(self.last_change_tick, self.change_tick)
+                .get_mut_using_ticks::<T>(self.last_change_tick, self.change_tick)
                 .ok_or(QueryComponentError::MissingComponent)
         } else {
             Err(QueryComponentError::MissingWriteAccess)

crates/bevy_ecs/src/system/system_param.rs

@@ -541,7 +541,8 @@ unsafe impl<'a, T: Resource> SystemParam for ResMut<'a, T> {
         change_tick: u32,
     ) -> Self::Item<'w, 's> {
         let value = world
-            .get_resource_unchecked_mut_with_id(component_id)
+            .as_unsafe_world_cell_migration_internal()
+            .get_resource_mut_with_id(component_id)
             .unwrap_or_else(|| {
                 panic!(
                     "Resource requested by {} does not exist: {}",
@@ -578,7 +579,8 @@ unsafe impl<'a, T: Resource> SystemParam for Option<ResMut<'a, T>> {
         change_tick: u32,
     ) -> Self::Item<'w, 's> {
         world
-            .get_resource_unchecked_mut_with_id(component_id)
+            .as_unsafe_world_cell_migration_internal()
+            .get_resource_mut_with_id(component_id)
             .map(|value| ResMut {
                 value: value.value,
                 ticks: TicksMut {