move take_component to Storages and add some safety comments

Author: jakobhellermann

crates/bevy_ecs/src/storage/mod.rs

@@ -10,11 +10,11 @@ pub use sparse_set::*;
 pub use table::*;
 
 use crate::{
-    archetype::Archetypes,
+    archetype::{Archetype, Archetypes},
     component::{ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entity, EntityLocation},
 };
-use bevy_ptr::Ptr;
+use bevy_ptr::{OwningPtr, Ptr};
 use std::{any::TypeId, cell::UnsafeCell};
 
 /// The raw data stores of a [World](crate::world::World)
@@ -193,3 +193,49 @@ impl Storages {
         }
     }
 }
+
+impl Storages {
+    /// Moves component data out of storage.
+    ///
+    /// This function leaves the underlying memory unchanged, but the component behind
+    /// returned pointer is semantically owned by the caller and will not be dropped in its original location.
+    /// Caller is responsible to drop component data behind returned pointer.
+    ///
+    /// # Safety
+    /// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside the `archetype`
+    /// - `component_id` must be valid
+    /// - `components` must come from the same world as `self`
+    /// - The relevant table row **must be removed** by the caller once all components are taken
+    #[inline]
+    pub(crate) unsafe fn take_component<'a>(
+        &'a mut self,
+        components: &Components,
+        archetype: &Archetype,
+        removed_components: &mut SparseSet<ComponentId, Vec<Entity>>,
+        component_id: ComponentId,
+        entity: Entity,
+        location: EntityLocation,
+    ) -> OwningPtr<'a> {
+        // SAFETY: component id is valid as per the functions safety requirements
+        let component_info = components.get_info_unchecked(component_id);
+        let removed_components = removed_components.get_or_insert_with(component_id, Vec::new);
+        removed_components.push(entity);
+        match component_info.storage_type() {
+            StorageType::Table => {
+                let table = &mut self.tables[archetype.table_id()];
+                // SAFETY: archetypes will always point to valid columns
+                let components = table.get_column_mut(component_id).unwrap();
+                let table_row = archetype.entity_table_row(location.index);
+                // SAFETY: archetypes only store valid table_rows and the stored component type is T
+                // SAFETY: promote is safe because the caller promises to remove the table row without dropping it immediately afterwards
+                components.get_data_unchecked_mut(table_row).promote()
+            }
+            StorageType::SparseSet => self
+                .sparse_sets
+                .get_mut(component_id)
+                .unwrap()
+                .remove_and_forget(entity)
+                .unwrap(),
+        }
+    }
+}

crates/bevy_ecs/src/world/entity_ref.rs

@@ -4,10 +4,10 @@ use crate::{
     change_detection::{MutUntyped, Ticks},
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entities, Entity, EntityLocation},
-    storage::{SparseSet, Storages},
+    storage::Storages,
     world::{Mut, World},
 };
-use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
+use bevy_ptr::{Ptr, UnsafeCellDeref};
 use bevy_utils::tracing::debug;
 use std::any::TypeId;
 
@@ -390,10 +390,12 @@ impl<'w> EntityMut<'w> {
         let result = unsafe {
             T::from_components(storages, &mut |storages| {
                 let component_id = bundle_components.next().unwrap();
-                // SAFETY: entity location is valid and table row is removed below
-                take_component(
+                // SAFETY:
+                // - entity location is valid
+                // - table row is removed below
+                // - `components` comes from the same world as `storages`
+                storages.take_component(
                     components,
-                    storages,
                     old_archetype,
                     removed_components,
                     component_id,
@@ -671,48 +673,6 @@ impl<'w> EntityMut<'w> {
     }
 }
 
-// TODO: move to Storages?
-/// Moves component data out of storage.
-///
-/// This function leaves the underlying memory unchanged, but the component behind
-/// returned pointer is semantically owned by the caller and will not be dropped in its original location.
-/// Caller is responsible to drop component data behind returned pointer.
-///
-/// # Safety
-/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside the archetype
-/// - `component_id` must be valid
-/// - The relevant table row **must be removed** by the caller once all components are taken
-#[inline]
-unsafe fn take_component<'a>(
-    components: &Components,
-    storages: &'a mut Storages,
-    archetype: &Archetype,
-    removed_components: &mut SparseSet<ComponentId, Vec<Entity>>,
-    component_id: ComponentId,
-    entity: Entity,
-    location: EntityLocation,
-) -> OwningPtr<'a> {
-    let component_info = components.get_info_unchecked(component_id);
-    let removed_components = removed_components.get_or_insert_with(component_id, Vec::new);
-    removed_components.push(entity);
-    match component_info.storage_type() {
-        StorageType::Table => {
-            let table = &mut storages.tables[archetype.table_id()];
-            // SAFETY: archetypes will always point to valid columns
-            let components = table.get_column_mut(component_id).unwrap();
-            let table_row = archetype.entity_table_row(location.index);
-            // SAFETY: archetypes only store valid table_rows and the stored component type is T
-            components.get_data_unchecked_mut(table_row).promote()
-        }
-        StorageType::SparseSet => storages
-            .sparse_sets
-            .get_mut(component_id)
-            .unwrap()
-            .remove_and_forget(entity)
-            .unwrap(),
-    }
-}
-
 fn contains_component_with_type(world: &World, type_id: TypeId, location: EntityLocation) -> bool {
     if let Some(component_id) = world.components.get_id(type_id) {
         contains_component_with_id(world, component_id, location)