move unchecked methods from world to InteriorMutableWorld

jakobhellermann · jakobhellermann · commit 96663daf2734 · 2022-12-17T13:18:42.000+01:00
diff --git a/crates/bevy_asset/src/reflect.rs b/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()
                 };
diff --git a/crates/bevy_ecs/src/reflect.rs b/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| {
diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs
@@ -521,7 +521,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: {}",
@@ -569,7 +570,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 {
diff --git a/crates/bevy_ecs/src/world/interior_mutable_world.rs b/crates/bevy_ecs/src/world/interior_mutable_world.rs
@@ -156,7 +156,7 @@ impl<'w> InteriorMutableWorld<'w> {
         // - component_id is of type `R`
         // - caller ensures aliasing rules
         // - `R` is Send + Sync
-        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.
@@ -195,19 +195,52 @@ impl<'w> InteriorMutableWorld<'w> {
             ticks,
         })
     }
+}
 
-    /// Gets a reference to the non-send resource of the given type, if it exists.
-    /// Otherwise returns [None]
+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
+    ///
+    /// All [`InteriorMutableWorld`] methods take `&self` and thus do not check that there is only one unique reference or multiple shared ones.
+    /// It is the callers responsibility to make sure that there will never be a mutable reference to a value that has other references pointing to it,
+    /// and that no arbitrary safe code can access a `&World` while some value is mutably borrowed.
     #[inline]
-    pub unsafe fn get_non_send_resource<R: 'static>(&self) -> Option<&R> {
-        self.0.get_non_send_resource::<R>()
+    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)?;
+
+        // SAFETY:
+        // - This caller ensures that nothing aliases `ticks`.
+        // - index is in-bounds because the column is initialized and non-empty
+        let ticks = unsafe {
+            Ticks::from_tick_cells(ticks, self.0.last_change_tick(), self.0.read_change_tick())
+        };
+
+        Some(Mut {
+            // SAFETY: caller ensures aliasing rules, ptr is of type `R`
+            value: unsafe { ptr.assert_unique().deref_mut() },
+            ticks,
+        })
     }
-    /// Gets a mutable reference to the non-send resource of the given type, if it exists.
-    /// Otherwise returns [None]
+
+    /// # 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.
+    ///
+    /// All [`InteriorMutableWorld`] methods take `&self` and thus do not check that there is only one unique reference or multiple shared ones.
+    /// It is the callers responsibility to make sure that there will never be a mutable reference to a value that has other references pointing to it,
+    /// and that no arbitrary safe code can access a `&World` while some value is mutably borrowed.
     #[inline]
-    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) }
+    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) }
     }
 }
diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs
@@ -913,7 +913,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
@@ -930,18 +930,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
@@ -995,20 +983,16 @@ impl World {
     /// Otherwise returns [None]
     #[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>> {
+        self.validate_non_send_access::<R>();
         let component_id = self.components.get_resource_id(TypeId::of::<R>())?;
-        self.get_non_send_unchecked_mut_with_id(component_id)
+        // SAFETY:
+        // - `component_id` is type `R`
+        // - we have unique access
+        // - main thread is validated
+        unsafe {
+            self.as_interior_mutable()
+                .get_resource_mut_with_id(component_id)
+        }
     }
 
     // Shorthand helper function for getting the data and change ticks for a resource.
@@ -1278,21 +1262,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::from_tick_cells(ticks, self.last_change_tick(), self.read_change_tick()),
-        })
-    }
-
     /// # Safety
     /// `component_id` must be assigned to a component of type `R`
     #[inline]
@@ -1304,18 +1273,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
@@ -1470,25 +1427,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 change_tick = self.change_tick();
-
-        let (ptr, ticks) = self.get_resource_with_ticks(component_id)?;
-
-        // SAFETY: This function has exclusive access to the world so nothing aliases `ticks`.
-        // - 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
-        let ticks = unsafe { Ticks::from_tick_cells(ticks, self.last_change_tick(), 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].
diff --git a/crates/bevy_ecs/src/world/world_cell.rs b/crates/bevy_ecs/src/world/world_cell.rs
@@ -226,7 +226,11 @@ impl<'w> WorldCell<'w> {
             .get_resource_archetype_component_id(component_id)?;
         WorldBorrowMut::try_new(
             // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
-            || unsafe { self.world.get_resource_unchecked_mut_with_id(component_id) },
+            || unsafe {
+                self.world
+                    .as_interior_mutable()
+                    .get_resource_mut_with_id(component_id)
+            },
             archetype_component_id,
             self.access.clone(),
         )
@@ -292,7 +296,11 @@ impl<'w> WorldCell<'w> {
             .get_resource_archetype_component_id(component_id)?;
         WorldBorrowMut::try_new(
             // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
-            || unsafe { self.world.get_non_send_unchecked_mut_with_id(component_id) },
+            || unsafe {
+                self.world
+                    .as_interior_mutable()
+                    .get_non_send_mut_with_id(component_id)
+            },
             archetype_component_id,
             self.access.clone(),
         )
