migrate WorldCell to InteriorMutableWorld

Author: jakobhellermann

crates/bevy_ecs/src/world/mod.rs

@@ -26,6 +26,7 @@ use bevy_ptr::{OwningPtr, Ptr};
 use bevy_utils::tracing::warn;
 use std::{
     any::TypeId,
+    cell::UnsafeCell,
     fmt,
     sync::atomic::{AtomicU32, Ordering},
 };
@@ -62,8 +63,8 @@ pub struct World {
     pub(crate) storages: Storages,
     pub(crate) bundles: Bundles,
     pub(crate) removed_components: SparseSet<ComponentId, Vec<Entity>>,
-    /// Access cache used by [WorldCell].
-    pub(crate) archetype_component_access: ArchetypeComponentAccess,
+    /// Access cache used by [WorldCell]. Is only accessed in the `Drop` impl of `WorldCell`.
+    pub(crate) archetype_component_access: UnsafeCell<ArchetypeComponentAccess>,
     pub(crate) change_tick: AtomicU32,
     pub(crate) last_change_tick: u32,
 }
@@ -1096,11 +1097,7 @@ impl World {
     #[inline]
     pub fn get_non_send_resource_mut<R: 'static>(&mut self) -> Option<Mut<'_, R>> {
         // SAFETY: unique world access
-        let component_id = self.components.get_resource_id(TypeId::of::<R>())?;
-        unsafe {
-            self.as_interior_mutable()
-                .get_non_send_mut_with_id::<R>(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.

crates/bevy_ecs/src/world/world_cell.rs

@@ -14,10 +14,12 @@ use std::{
     rc::Rc,
 };
 
+use super::interior_mutable_world::InteriorMutableWorld;
+
 /// Exposes safe mutable access to multiple resources at a time in a World. Attempting to access
 /// World in a way that violates Rust's mutability rules will panic thanks to runtime checks.
 pub struct WorldCell<'w> {
-    pub(crate) world: &'w mut World,
+    pub(crate) world: InteriorMutableWorld<'w>,
     pub(crate) access: Rc<RefCell<ArchetypeComponentAccess>>,
 }
 
@@ -76,8 +78,16 @@ impl ArchetypeComponentAccess {
 impl<'w> Drop for WorldCell<'w> {
     fn drop(&mut self) {
         let mut access = self.access.borrow_mut();
-        // give world ArchetypeComponentAccess back to reuse allocations
-        std::mem::swap(&mut self.world.archetype_component_access, &mut *access);
+
+        {
+            // SAFETY: we only swap `archetype_component_access`
+            let world = unsafe { self.world.world() };
+            // SAFETY: the WorldCell has exclusive world access
+            let world_cached_access = unsafe { &mut *world.archetype_component_access.get() };
+
+            // give world ArchetypeComponentAccess back to reuse allocations
+            std::mem::swap(world_cached_access, &mut *access);
+        }
     }
 }
 
@@ -175,25 +185,25 @@ impl<'w> WorldCell<'w> {
     pub(crate) fn new(world: &'w mut World) -> Self {
         // this is cheap because ArchetypeComponentAccess::new() is const / allocation free
         let access = std::mem::replace(
-            &mut world.archetype_component_access,
+            world.archetype_component_access.get_mut(),
             ArchetypeComponentAccess::new(),
         );
         // world's ArchetypeComponentAccess is recycled to cut down on allocations
         Self {
-            world,
+            world: world.as_interior_mutable(),
             access: Rc::new(RefCell::new(access)),
         }
     }
 
     /// Gets a reference to the resource of the given type
     pub fn get_resource<T: Resource>(&self) -> Option<WorldBorrow<'_, T>> {
-        let component_id = self.world.components.get_resource_id(TypeId::of::<T>())?;
-        let archetype_component_id = self
-            .world
-            .get_resource_archetype_component_id(component_id)?;
+        let component_id = self.world.components().get_resource_id(TypeId::of::<T>())?;
+
+        let archetype_component_id = self.world.storages().resources.get(component_id)?.id();
+
         WorldBorrow::try_new(
-            // SAFETY: ComponentId matches TypeId
-            || unsafe { self.world.get_resource_with_id(component_id) },
+            // SAFETY: access is checked by WorldBorrow
+            || unsafe { self.world.get_resource::<T>() },
             archetype_component_id,
             self.access.clone(),
         )
@@ -220,17 +230,11 @@ impl<'w> WorldCell<'w> {
 
     /// Gets a mutable reference to the resource of the given type
     pub fn get_resource_mut<T: Resource>(&self) -> Option<WorldBorrowMut<'_, T>> {
-        let component_id = self.world.components.get_resource_id(TypeId::of::<T>())?;
-        let archetype_component_id = self
-            .world
-            .get_resource_archetype_component_id(component_id)?;
+        let component_id = self.world.components().get_resource_id(TypeId::of::<T>())?;
+        let archetype_component_id = self.world.storages().resources.get(component_id)?.id();
         WorldBorrowMut::try_new(
-            // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
-            || unsafe {
-                self.world
-                    .as_interior_mutable_migration_internal()
-                    .get_resource_mut_with_id(component_id)
-            },
+            // SAFETY: access is checked by WorldBorrowMut
+            || unsafe { self.world.get_resource_mut::<T>() },
             archetype_component_id,
             self.access.clone(),
         )
@@ -257,13 +261,16 @@ impl<'w> WorldCell<'w> {
 
     /// Gets an immutable reference to the non-send resource of the given type, if it exists.
     pub fn get_non_send_resource<T: 'static>(&self) -> Option<WorldBorrow<'_, T>> {
-        let component_id = self.world.components.get_resource_id(TypeId::of::<T>())?;
+        let component_id = self.world.components().get_resource_id(TypeId::of::<T>())?;
         let archetype_component_id = self
             .world
-            .get_non_send_archetype_component_id(component_id)?;
+            .storages()
+            .non_send_resources
+            .get(component_id)?
+            .id();
         WorldBorrow::try_new(
-            // SAFETY: ComponentId matches TypeId
-            || unsafe { self.world.get_non_send_with_id(component_id) },
+            // SAFETY: access is checked by WorldBorrowMut
+            || unsafe { self.world.get_non_send_resource::<T>() },
             archetype_component_id,
             self.access.clone(),
         )
@@ -290,17 +297,16 @@ impl<'w> WorldCell<'w> {
 
     /// Gets a mutable reference to the non-send resource of the given type, if it exists.
     pub fn get_non_send_resource_mut<T: 'static>(&self) -> Option<WorldBorrowMut<'_, T>> {
-        let component_id = self.world.components.get_resource_id(TypeId::of::<T>())?;
+        let component_id = self.world.components().get_resource_id(TypeId::of::<T>())?;
         let archetype_component_id = self
             .world
-            .get_non_send_archetype_component_id(component_id)?;
+            .storages()
+            .non_send_resources
+            .get(component_id)?
+            .id();
         WorldBorrowMut::try_new(
             // SAFETY: access is checked by WorldBorrowMut
-            || unsafe {
-                self.world
-                    .as_interior_mutable_migration_internal()
-                    .get_non_send_resource_mut::<T>()
-            },
+            || unsafe { self.world.get_non_send_resource_mut::<T>() },
             archetype_component_id,
             self.access.clone(),
         )
@@ -416,10 +422,11 @@ mod tests {
         let u32_archetype_component_id = world
             .get_resource_archetype_component_id(u32_component_id)
             .unwrap();
-        assert_eq!(world.archetype_component_access.access.len(), 1);
+        assert_eq!(world.archetype_component_access.get_mut().access.len(), 1);
         assert_eq!(
             world
                 .archetype_component_access
+                .get_mut()
                 .access
                 .get(u32_archetype_component_id),
             Some(&BASE_ACCESS),