Provide access to world storages via UnsafeWorldCell (#8987)

# Objective Title. This is necessary in order to update [`bevy-trait-query`](https://crates.io/crates/bevy-trait-query) to Bevy 0.11. --- ## Changelog Added the unsafe function `UnsafeWorldCell::storages`, which provides unchecked access to the internal data stores of a `World`.
bevyengine · Jun 29, 2023 · de1dcb9 · de1dcb9
1 parent e29981d
commit de1dcb9
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 29 deletions.
diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs
@@ -662,7 +662,6 @@ unsafe impl<T: Component> WorldQuery for &T {
                     // which we are allowed to access since we registered it in `update_archetype_component_access`.
                     // Note that we do not actually access any components in this function, we just get a shared
                     // reference to the sparse set, which is used to access the components in `Self::fetch`.
-                    .unsafe_world()
                     .storages()
                     .sparse_sets
                     .get(component_id)
@@ -810,7 +809,6 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
             sparse_set: (T::Storage::STORAGE_TYPE == StorageType::SparseSet).then(|| {
                 world
                     // SAFETY: See &T::init_fetch.
-                    .unsafe_world()
                     .storages()
                     .sparse_sets
                     .get(component_id)
@@ -974,7 +972,6 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
             sparse_set: (T::Storage::STORAGE_TYPE == StorageType::SparseSet).then(|| {
                 world
                     // SAFETY: See &T::init_fetch.
-                    .unsafe_world()
                     .storages()
                     .sparse_sets
                     .get(component_id)

diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs
@@ -34,7 +34,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> QueryIter<'w, 's, Q, F> {
         QueryIter {
             query_state,
             // SAFETY: We only access table data that has been registered in `query_state`.
-            tables: &world.unsafe_world().storages().tables,
+            tables: &world.storages().tables,
             archetypes: world.archetypes(),
             cursor: QueryIterationCursor::init(world, query_state, last_run, this_run),
         }
@@ -107,7 +107,7 @@ where
             archetypes: world.archetypes(),
             // SAFETY: We only access table data that has been registered in `query_state`.
             // This means `world` has permission to access the data we use.
-            tables: &world.unsafe_world().storages.tables,
+            tables: &world.storages().tables,
             fetch,
             filter,
             entity_iter: entity_list.into_iter(),
@@ -316,7 +316,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery, const K: usize>
         QueryCombinationIter {
             query_state,
             // SAFETY: We only access table data that has been registered in `query_state`.
-            tables: &world.unsafe_world().storages().tables,
+            tables: &world.storages().tables,
             archetypes: world.archetypes(),
             cursors: array.assume_init(),
         }

diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs
@@ -462,7 +462,6 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         let mut filter = F::init_fetch(world, &self.filter_state, last_run, this_run);
 
         let table = world
-            .unsafe_world()
             .storages()
             .tables
             .get(location.table_id)
@@ -973,7 +972,7 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         let mut fetch = Q::init_fetch(world, &self.fetch_state, last_run, this_run);
         let mut filter = F::init_fetch(world, &self.filter_state, last_run, this_run);
 
-        let tables = &world.unsafe_world().storages().tables;
+        let tables = &world.storages().tables;
         if Q::IS_DENSE && F::IS_DENSE {
             for table_id in &self.matched_table_ids {
                 let table = tables.get(*table_id).debug_checked_unwrap();
@@ -1048,7 +1047,7 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         ComputeTaskPool::get().scope(|scope| {
             if Q::IS_DENSE && F::IS_DENSE {
                 // SAFETY: We only access table data that has been registered in `self.archetype_component_access`.
-                let tables = &world.unsafe_world().storages().tables;
+                let tables = &world.storages().tables;
                 for table_id in &self.matched_table_ids {
                     let table = &tables[*table_id];
                     if table.is_empty() {
@@ -1064,7 +1063,7 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
                                 Q::init_fetch(world, &self.fetch_state, last_run, this_run);
                             let mut filter =
                                 F::init_fetch(world, &self.filter_state, last_run, this_run);
-                            let tables = &world.unsafe_world().storages().tables;
+                            let tables = &world.storages().tables;
                             let table = tables.get(*table_id).debug_checked_unwrap();
                             let entities = table.entities();
                             Q::set_table(&mut fetch, &self.fetch_state, table);
@@ -1108,7 +1107,7 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
                                 Q::init_fetch(world, &self.fetch_state, last_run, this_run);
                             let mut filter =
                                 F::init_fetch(world, &self.filter_state, last_run, this_run);
-                            let tables = &world.unsafe_world().storages().tables;
+                            let tables = &world.storages().tables;
                             let archetype =
                                 world.archetypes().get(*archetype_id).debug_checked_unwrap();
                             let table = tables.get(archetype.table_id()).debug_checked_unwrap();

diff --git a/crates/bevy_ecs/src/world/unsafe_world_cell.rs b/crates/bevy_ecs/src/world/unsafe_world_cell.rs
@@ -12,7 +12,7 @@ use crate::{
     },
     entity::{Entities, Entity, EntityLocation},
     prelude::Component,
-    storage::{Column, ComponentSparseSet},
+    storage::{Column, ComponentSparseSet, Storages},
     system::Resource,
 };
 use bevy_ptr::Ptr;
@@ -271,6 +271,20 @@ impl<'w> UnsafeWorldCell<'w> {
         unsafe { self.world_metadata() }.increment_change_tick()
     }
 
+    /// Provides unchecked access to the internal data stores of the [`World`].
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that this is only used to access world data
+    /// that this [`UnsafeWorldCell`] is allowed to.
+    /// As always, any mutable access to a component must not exist at the same
+    /// time as any other accesses to that same component.
+    #[inline]
+    pub unsafe fn storages(self) -> &'w Storages {
+        // SAFETY: The caller promises to only access world data allowed by this instance.
+        &unsafe { self.unsafe_world() }.storages
+    }
+
     /// Shorthand helper function for getting the [`ArchetypeComponentId`] for a resource.
     #[inline]
     pub(crate) fn get_resource_archetype_component_id(
@@ -342,8 +356,7 @@ impl<'w> UnsafeWorldCell<'w> {
     pub unsafe fn get_resource_by_id(self, component_id: ComponentId) -> Option<Ptr<'w>> {
         // SAFETY: caller ensures that `self` has permission to access `R`
         //  caller ensures that no mutable reference exists to `R`
-        unsafe { self.unsafe_world() }
-            .storages
+        unsafe { self.storages() }
             .resources
             .get(component_id)?
             .get_data()
@@ -385,8 +398,7 @@ impl<'w> UnsafeWorldCell<'w> {
     pub unsafe fn get_non_send_resource_by_id(self, component_id: ComponentId) -> Option<Ptr<'w>> {
         // SAFETY: we only access data on world that the caller has ensured is unaliased and we have
         //  permission to access.
-        unsafe { self.unsafe_world() }
-            .storages
+        unsafe { self.storages() }
             .non_send_resources
             .get(component_id)?
             .get_data()
@@ -429,8 +441,7 @@ impl<'w> UnsafeWorldCell<'w> {
     ) -> Option<MutUntyped<'w>> {
         // SAFETY: we only access data that the caller has ensured is unaliased and `self`
         //  has permission to access.
-        let (ptr, ticks) = unsafe { self.unsafe_world() }
-            .storages
+        let (ptr, ticks) = unsafe { self.storages() }
             .resources
             .get(component_id)?
             .get_with_ticks()?;
@@ -492,8 +503,7 @@ impl<'w> UnsafeWorldCell<'w> {
         let change_tick = self.change_tick();
         // SAFETY: we only access data that the caller has ensured is unaliased and `self`
         //  has permission to access.
-        let (ptr, ticks) = unsafe { self.unsafe_world() }
-            .storages
+        let (ptr, ticks) = unsafe { self.storages() }
             .non_send_resources
             .get(component_id)?
             .get_with_ticks()?;
@@ -526,8 +536,7 @@ impl<'w> UnsafeWorldCell<'w> {
         // - caller ensures there is no `&mut World`
         // - caller ensures there are no mutable borrows of this resource
         // - caller ensures that we have permission to access this resource
-        unsafe { self.unsafe_world() }
-            .storages
+        unsafe { self.storages() }
             .resources
             .get(component_id)?
             .get_with_ticks()
@@ -551,8 +560,7 @@ impl<'w> UnsafeWorldCell<'w> {
         // - caller ensures there is no `&mut World`
         // - caller ensures there are no mutable borrows of this resource
         // - caller ensures that we have permission to access this resource
-        unsafe { self.unsafe_world() }
-            .storages
+        unsafe { self.storages() }
             .non_send_resources
             .get(component_id)?
             .get_with_ticks()
@@ -878,7 +886,7 @@ impl<'w> UnsafeWorldCell<'w> {
     ) -> Option<&'w Column> {
         // SAFETY: caller ensures returned data is not misused and we have not created any borrows
         // of component/resource data
-        unsafe { self.unsafe_world() }.storages.tables[location.table_id].get_column(component_id)
+        unsafe { self.storages() }.tables[location.table_id].get_column(component_id)
     }
 
     #[inline]
@@ -889,10 +897,7 @@ impl<'w> UnsafeWorldCell<'w> {
     unsafe fn fetch_sparse_set(self, component_id: ComponentId) -> Option<&'w ComponentSparseSet> {
         // SAFETY: caller ensures returned data is not misused and we have not created any borrows
         // of component/resource data
-        unsafe { self.unsafe_world() }
-            .storages
-            .sparse_sets
-            .get(component_id)
+        unsafe { self.storages() }.sparse_sets.get(component_id)
     }
 }