move unchecked methods from world to InteriorMutableWorld

Author: jakobhellermann

crates/bevy_asset/src/reflect.rs

@@ -148,7 +148,8 @@ impl<A: Asset + FromReflect> FromType<A> for ReflectAsset {
             get_unchecked_mut: |world, handle| {
                 let assets = unsafe {
                     world
-                        .get_resource_unchecked_mut::<Assets<A>>()
+                        .as_interior_mutable()
+                        .get_resource_mut::<Assets<A>>()
                         .unwrap()
                         .into_inner()
                 };

crates/bevy_ecs/src/reflect.rs

@@ -385,10 +385,13 @@ 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 {
-                        value: res.value as &mut dyn Reflect,
-                        ticks: res.ticks,
-                    })
+                    world
+                        .as_interior_mutable()
+                        .get_resource_mut::<C>()
+                        .map(|res| Mut {
+                            value: res.value as &mut dyn Reflect,
+                            ticks: res.ticks,
+                        })
                 }
             },
             copy: |source_world, destination_world| {

crates/bevy_ecs/src/system/system_param.rs

@@ -515,7 +515,8 @@ impl<'w, 's, T: Resource> SystemParamFetch<'w, 's> for ResMutState<T> {
         change_tick: u32,
     ) -> Self::Item {
         let value = world
-            .get_resource_unchecked_mut_with_id(state.component_id)
+            .as_interior_mutable()
+            .get_resource_mut_with_id(state.component_id)
             .unwrap_or_else(|| {
                 panic!(
                     "Resource requested by {} does not exist: {}",
@@ -562,7 +563,8 @@ impl<'w, 's, T: Resource> SystemParamFetch<'w, 's> for OptionResMutState<T> {
         change_tick: u32,
     ) -> Self::Item {
         world
-            .get_resource_unchecked_mut_with_id(state.0.component_id)
+            .as_interior_mutable()
+            .get_resource_mut_with_id(state.0.component_id)
             .map(|value| ResMut {
                 value: value.value,
                 ticks: Ticks {

crates/bevy_ecs/src/world/interior_mutable_world.rs

@@ -102,7 +102,7 @@ impl<'w> InteriorMutableWorld<'w> {
     #[inline]
     pub unsafe fn get_resource_mut<R: Resource>(&self) -> Option<Mut<'w, R>> {
         let component_id = self.0.components.get_resource_id(TypeId::of::<R>())?;
-        unsafe { self.0.get_resource_unchecked_mut_with_id(component_id) }
+        unsafe { self.get_resource_mut_with_id(component_id) }
     }
 
     /// Gets a pointer to the resource with the id [`ComponentId`] if it exists.
@@ -152,6 +152,43 @@ impl<'w> InteriorMutableWorld<'w> {
     pub unsafe fn get_non_send_resource_mut<R: 'static>(&mut self) -> Option<Mut<'w, R>> {
         let component_id = self.0.components.get_resource_id(TypeId::of::<R>())?;
         // SAFETY: safety requirement is deferred to the caller
-        unsafe { self.0.get_non_send_unchecked_mut_with_id(component_id) }
+        unsafe { self.get_non_send_mut_with_id(component_id) }
+    }
+}
+
+impl<'w> InteriorMutableWorld<'w> {
+    /// # Safety
+    /// - `component_id` must be assigned to a component of type `R`
+    /// - Caller must ensure this doesn't violate Rust mutability rules for the given resource.
+    /// - resource is either Send+Sync or the main thread is validated
+    #[inline]
+    pub(crate) unsafe fn get_resource_mut_with_id<R>(
+        &self,
+        component_id: ComponentId,
+    ) -> Option<Mut<'w, R>> {
+        let (ptr, ticks) = self.0.get_resource_with_ticks(component_id)?;
+        Some(Mut {
+            // caller ensures aliasing rules, ptr is of type `R`
+            value: unsafe { ptr.assert_unique().deref_mut() },
+            ticks: Ticks {
+                // caller ensures aliasing rules
+                component_ticks: unsafe { ticks.deref_mut() },
+                last_change_tick: self.0.last_change_tick(),
+                change_tick: self.0.read_change_tick(),
+            },
+        })
+    }
+
+    /// # Safety
+    /// - `component_id` must be assigned to a component of type `R`.
+    /// - Caller must ensure this doesn't violate Rust mutability rules for the given resource.
+    #[inline]
+    pub(crate) unsafe fn get_non_send_mut_with_id<R: 'static>(
+        &self,
+        component_id: ComponentId,
+    ) -> Option<Mut<'w, R>> {
+        self.0.validate_non_send_access::<R>();
+        // SAFETY: we validated nonsend access, the rest of the requirements are deferred to caller
+        unsafe { self.get_resource_mut_with_id(component_id) }
     }
 }

crates/bevy_ecs/src/world/mod.rs

@@ -20,7 +20,7 @@ use crate::{
     storage::{ResourceData, SparseSet, Storages},
     system::Resource,
 };
-use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
+use bevy_ptr::{OwningPtr, Ptr};
 use bevy_utils::tracing::warn;
 use std::{
     any::TypeId,
@@ -907,7 +907,7 @@ impl World {
     #[inline]
     pub fn get_resource_mut<R: Resource>(&mut self) -> Option<Mut<'_, R>> {
         // SAFETY: unique world access
-        unsafe { self.get_resource_unchecked_mut() }
+        unsafe { self.as_interior_mutable().get_resource_mut() }
     }
 
     // PERF: optimize this to avoid redundant lookups
@@ -924,18 +924,6 @@ impl World {
         self.resource_mut()
     }
 
-    /// Gets a mutable reference to the resource of the given type, if it exists
-    /// Otherwise returns [None]
-    ///
-    /// # Safety
-    /// This will allow aliased mutable access to the given resource type. The caller must ensure
-    /// that there is either only one mutable access or multiple immutable accesses at a time.
-    #[inline]
-    pub unsafe fn get_resource_unchecked_mut<R: Resource>(&self) -> Option<Mut<'_, R>> {
-        let component_id = self.components.get_resource_id(TypeId::of::<R>())?;
-        self.get_resource_unchecked_mut_with_id(component_id)
-    }
-
     /// Gets an immutable reference to the non-send resource of the given type, if it exists.
     ///
     /// # Panics
@@ -990,19 +978,7 @@ impl World {
     #[inline]
     pub fn get_non_send_resource_mut<R: 'static>(&mut self) -> Option<Mut<'_, R>> {
         // SAFETY: unique world access
-        unsafe { self.get_non_send_resource_unchecked_mut() }
-    }
-
-    /// Gets a mutable reference to the non-send resource of the given type, if it exists.
-    /// Otherwise returns [None]
-    ///
-    /// # Safety
-    /// This will allow aliased mutable access to the given non-send resource type. The caller must
-    /// ensure that there is either only one mutable access or multiple immutable accesses at a time.
-    #[inline]
-    pub unsafe fn get_non_send_resource_unchecked_mut<R: 'static>(&self) -> Option<Mut<'_, R>> {
-        let component_id = self.components.get_resource_id(TypeId::of::<R>())?;
-        self.get_non_send_unchecked_mut_with_id(component_id)
+        unsafe { self.as_interior_mutable().get_non_send_resource_mut() }
     }
 
     // Shorthand helper function for getting the data and change ticks for a resource.
@@ -1271,25 +1247,6 @@ impl World {
             .map(|ptr| ptr.deref())
     }
 
-    /// # Safety
-    /// `component_id` must be assigned to a component of type `R`
-    /// Caller must ensure this doesn't violate Rust mutability rules for the given resource.
-    #[inline]
-    pub(crate) unsafe fn get_resource_unchecked_mut_with_id<R>(
-        &self,
-        component_id: ComponentId,
-    ) -> Option<Mut<'_, R>> {
-        let (ptr, ticks) = self.get_resource_with_ticks(component_id)?;
-        Some(Mut {
-            value: ptr.assert_unique().deref_mut(),
-            ticks: Ticks {
-                component_ticks: ticks.deref_mut(),
-                last_change_tick: self.last_change_tick(),
-                change_tick: self.read_change_tick(),
-            },
-        })
-    }
-
     /// # Safety
     /// `component_id` must be assigned to a component of type `R`
     #[inline]
@@ -1301,18 +1258,6 @@ impl World {
         self.get_resource_with_id(component_id)
     }
 
-    /// # Safety
-    /// `component_id` must be assigned to a component of type `R`.
-    /// Caller must ensure this doesn't violate Rust mutability rules for the given resource.
-    #[inline]
-    pub(crate) unsafe fn get_non_send_unchecked_mut_with_id<R: 'static>(
-        &self,
-        component_id: ComponentId,
-    ) -> Option<Mut<'_, R>> {
-        self.validate_non_send_access::<R>();
-        self.get_resource_unchecked_mut_with_id(component_id)
-    }
-
     /// Inserts a new resource with the given `value`. Will replace the value if it already existed.
     ///
     /// **You should prefer to use the typed API [`World::insert_resource`] where possible and only
@@ -1459,28 +1404,11 @@ impl World {
     /// use this in cases where the actual types are not known at compile time.**
     #[inline]
     pub fn get_resource_mut_by_id(&mut self, component_id: ComponentId) -> Option<MutUntyped<'_>> {
-        let info = self.components.get_info(component_id)?;
-        if !info.is_send_and_sync() {
-            self.validate_non_send_access_untyped(info.name());
+        // SAFETY: unique world access
+        unsafe {
+            self.as_interior_mutable()
+                .get_resource_mut_by_id(component_id)
         }
-
-        let (ptr, ticks) = self.get_resource_with_ticks(component_id)?;
-
-        // SAFE: This function has exclusive access to the world so nothing aliases `ticks`.
-        let ticks = Ticks {
-            // SAFETY:
-            // - index is in-bounds because the column is initialized and non-empty
-            // - no other reference to the ticks of the same row can exist at the same time
-            component_ticks: unsafe { ticks.deref_mut() },
-            last_change_tick: self.last_change_tick(),
-            change_tick: self.read_change_tick(),
-        };
-
-        Some(MutUntyped {
-            // SAFETY: This function has exclusive access to the world so nothing aliases `ptr`.
-            value: unsafe { ptr.assert_unique() },
-            ticks,
-        })
     }
 
     /// Removes the resource of a given type, if it exists. Otherwise returns [None].

crates/bevy_ecs/src/world/world_cell.rs

@@ -226,7 +226,8 @@ impl<'w> WorldCell<'w> {
             // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
             unsafe {
                 self.world
-                    .get_resource_unchecked_mut_with_id(component_id)?
+                    .as_interior_mutable()
+                    .get_resource_mut_with_id(component_id)?
             },
             archetype_component_id,
             self.access.clone(),
@@ -295,7 +296,8 @@ impl<'w> WorldCell<'w> {
             // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
             unsafe {
                 self.world
-                    .get_non_send_unchecked_mut_with_id(component_id)?
+                    .as_interior_mutable()
+                    .get_non_send_mut_with_id(component_id)?
             },
             archetype_component_id,
             self.access.clone(),