Add an API for accessing world storages from UnsafeWorldCell

crates/bevy_ecs/src/storage/mod.rs

@@ -29,6 +29,8 @@ pub use resource::*;
 pub use sparse_set::*;
 pub use table::*;
 
+use crate::component::ComponentId;
+
 /// The raw data stores of a [World](crate::world::World)
 #[derive(Default)]
 pub struct Storages {
@@ -41,3 +43,71 @@ pub struct Storages {
     /// Backing storage for `!Send` resources.
     pub non_send_resources: Resources<false>,
 }
+
+/// Provides interior-mutable access to a world's internal data storages.
+///
+/// Any instance of this type is associated with a set of world data that
+/// it is allowed to access. This should be described in the documentation
+/// of wherever you obtained the `UnsafeStorages`.
+///
+/// For instance, if you originally obtained it from a system running on
+/// a multi-threaded executor, then you are only allowed to access data
+/// that has been registered in the system's `archetype_component_access`.
+/// If you originally obtained an `UnsafeStorages` from an `&World`,
+/// then you have read-only access to the entire world.
+///
+/// Accessing world data that do not have access to, or mutably accessing
+/// data that you only have read-access to, is considered undefined behavior.
+pub struct UnsafeStorages<'a> {
+    pub sparse_sets: UnsafeSparseSets<'a>,
+    pub tables: UnsafeTables<'a>,
+    pub resources: UnsafeResources<'a, true>,
+    pub non_send_resources: UnsafeResources<'a, false>,
+}
+
+impl<'a> UnsafeStorages<'a> {
+    pub(crate) fn new(storages: &'a Storages) -> Self {
+        Self {
+            sparse_sets: UnsafeSparseSets {
+                sparse_sets: &storages.sparse_sets,
+            },
+            tables: UnsafeTables {
+                tables: &storages.tables,
+            },
+            resources: UnsafeResources {
+                resources: &storages.resources,
+            },
+            non_send_resources: UnsafeResources {
+                resources: &storages.non_send_resources,
+            },
+        }
+    }
+}
+
+/// A view into the [`ComponentSparseSet`] collection of [`UnsafeStorages`].
+#[derive(Clone, Copy)]
+pub struct UnsafeSparseSets<'a> {
+    sparse_sets: &'a SparseSets,
+}
+
+impl<'a> UnsafeSparseSets<'a> {
+    /// Gets a view into the [`ComponentSparseSet`] associated with `component_id`.
+    pub fn get(self, component_id: ComponentId) -> Option<UnsafeComponentSparseSet<'a>> {
+        self.sparse_sets
+            .get(component_id)
+            .map(|sparse_set| UnsafeComponentSparseSet { sparse_set })
+    }
+}
+
+/// A view into the [`Table`] collection of [`UnsafeStorages`].
+#[derive(Clone, Copy)]
+pub struct UnsafeTables<'a> {
+    tables: &'a Tables,
+}
+
+impl<'a> UnsafeTables<'a> {
+    /// Gets a view into the [`Table`] associated with `id`.
+    pub fn get(self, id: TableId) -> Option<UnsafeTable<'a>> {
+        self.tables.get(id).map(|table| UnsafeTable { table })
+    }
+}

crates/bevy_ecs/src/storage/resource.rs

@@ -286,3 +286,25 @@ impl<const SEND: bool> Resources<SEND> {
         }
     }
 }
+
+/// A view into a [`Resources`] collection of [`UnsafeStorages`].
+///
+/// [`UnsafeStorages`]: super::UnsafeStorages
+#[derive(Clone, Copy)]
+pub struct UnsafeResources<'a, const SEND: bool> {
+    pub(super) resources: &'a Resources<SEND>,
+}
+
+impl<'a, const SEND: bool> UnsafeResources<'a, SEND> {
+    /// Returns the entity's component and associated ticks.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeWorldCell`] self as obtained from has permission to access the resource mutably
+    /// - no mutable reference to the resource exists at the same time
+    ///
+    /// [`UnsafeWorldCell`]: crate::world::unsafe_world_cell::UnsafeWorldCell
+    pub unsafe fn get(self, component_id: ComponentId) -> Option<&'a ResourceData<SEND>> {
+        self.resources.get(component_id)
+    }
+}

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -626,6 +626,53 @@ impl SparseSets {
     }
 }
 
+/// A view into a [`ComponentSparseSet`] from [`UnsafeStorages`].
+///
+/// [`UnsafeStorages`]: super::UnsafeStorages
+#[derive(Clone, Copy)]
+pub struct UnsafeComponentSparseSet<'a> {
+    pub(super) sparse_set: &'a ComponentSparseSet,
+}
+
+impl<'a> UnsafeComponentSparseSet<'a> {
+    /// Returns the entity's component and associated ticks.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeWorldCell`] self was obtianed from has permission
+    ///   to access the component.
+    /// - no other mutable references to the component exist at the same time.
+    ///
+    /// [`UnsafeWorldCell`]: crate::world::unsafe_world_cell::UnsafeWorldCell
+    pub unsafe fn get(self, entity: Entity) -> Option<(Ptr<'a>, TickCells<'a>)> {
+        self.sparse_set.get_with_ticks(entity)
+    }
+
+    /// If `entity` has a component in this sparse set, returns the tick
+    /// indicating when it was added.
+    ///
+    /// Before dereferencing this, take care to ensure that the instance
+    /// of [`UnsafeStorages`] that this `UnsafeComponentSparseSet` was
+    /// obtained from has permission to access this component's data.
+    ///
+    /// [`UnsafeStorages`]: super::UnsafeStorages
+    pub fn get_added_tick(self, entity: Entity) -> Option<&'a UnsafeCell<Tick>> {
+        self.sparse_set.get_added_ticks(entity)
+    }
+
+    /// If `entity` has a component in this sparse set, returns the tick
+    /// indicating when it was last changed.
+    ///
+    /// Before dereferencing this, take care to ensure that the instance
+    /// of [`UnsafeStorages`] that this `UnsafeComponentSparseSet` was
+    /// obtained from has permission to access this component's data.
+    ///
+    /// [`UnsafeStorages`]: super::UnsafeStorages
+    pub fn get_changed_tick(self, entity: Entity) -> Option<&'a UnsafeCell<Tick>> {
+        self.sparse_set.get_changed_ticks(entity)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::SparseSets;

crates/bevy_ecs/src/storage/table.rs

@@ -559,7 +559,7 @@ pub struct Table {
 }
 
 impl Table {
-    /// Fetches a read-only slice of the entities stored within the [`Table`].
+    /// Returns the entities with components stored in this table.
     #[inline]
     pub fn entities(&self) -> &[Entity] {
         &self.entities
@@ -909,6 +909,44 @@ impl IndexMut<TableId> for Tables {
     }
 }
 
+/// A view into a [`Table`] from [`UnsafeStorages`].
+///
+/// [`UnsafeStorages`]: super::UnsafeStorages
+#[derive(Clone, Copy)]
+pub struct UnsafeTable<'a> {
+    pub(super) table: &'a Table,
+}
+
+impl<'a> UnsafeTable<'a> {
+    /// Gets access to the data for a specific component type in this table.
+    ///
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeWorldCell`] self was obtained from has permission
+    ///   to access the component in this table's archetype.
+    /// - no other mutable references to components of this type
+    ///   exist in this table at the same time.
+    ///
+    /// [`UnsafeWorldCell`]: crate::world::unsafe_world_cell::UnsafeWorldCell
+    pub unsafe fn get_column(self, component_id: ComponentId) -> Option<&'a Column> {
+        self.table.get_column(component_id)
+    }
+
+    /// Checks if the table contains a [`Column`] for a given [`Component`].
+    ///
+    /// Returns `true` if the column is present, `false` otherwise.
+    ///
+    /// [`Component`]: crate::component::Component
+    pub fn has_column(self, component_id: ComponentId) -> bool {
+        self.table.has_column(component_id)
+    }
+
+    /// Returns the entities with components stored in this table.
+    pub fn entities(self) -> &'a [Entity] {
+        self.table.entities()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate as bevy_ecs;

crates/bevy_ecs/src/world/unsafe_world_cell.rs

@@ -10,7 +10,7 @@ use crate::{
     },
     entity::{Entities, Entity, EntityLocation},
     prelude::Component,
-    storage::{Column, ComponentSparseSet},
+    storage::{Column, ComponentSparseSet, UnsafeStorages},
     system::Resource,
 };
 use bevy_ptr::Ptr;
@@ -151,6 +151,16 @@ impl<'w> UnsafeWorldCell<'w> {
         unsafe { &*self.0 }
     }
 
+    /// Returns interior-mutable access to the world's internal data stores.
+    /// The returned [`UnsafeStorages`] can only be used to access data
+    /// that this `UnsafeWorldCell` has permission to access.
+    pub fn storages(self) -> UnsafeStorages<'w> {
+        // SAFETY: Interior mutable access to the world is hidden behind
+        // `UnsafeStorages`, which will require any data access to be valid.
+        let storages = unsafe { self.unsafe_world().storages() };
+        UnsafeStorages::new(storages)
+    }
+
     /// Retrieves this world's [Entities] collection
     #[inline]
     pub fn entities(self) -> &'w Entities {