[Merged by Bors] - reduce tricky unsafety and simplify table structure

crates/bevy_ecs/src/bundle.rs

@@ -130,7 +130,7 @@ impl BundleInfo {
         &self,
         sparse_sets: &mut SparseSets,
         entity: Entity,
-        table: &Table,
+        table: &mut Table,
         table_row: usize,
         bundle_status: &[ComponentStatus],
         bundle: T,
@@ -145,8 +145,8 @@ impl BundleInfo {
             let component_status = bundle_status.get_unchecked(bundle_component);
             match self.storage_types[bundle_component] {
                 StorageType::Table => {
-                    let column = table.get_column(component_id).unwrap();
-                    column.set_unchecked(table_row, component_ptr);
+                    let column = table.get_column_mut(component_id).unwrap();
+                    column.set_data_unchecked(table_row, component_ptr);
                     let column_status = column.get_ticks_unchecked_mut(table_row);
                     match component_status {
                         ComponentStatus::Added => {

crates/bevy_ecs/src/storage/blob_vec.rs

@@ -35,7 +35,7 @@ impl BlobVec {
                 item_layout,
                 drop,
             };
-            blob_vec.reserve(capacity);
+            blob_vec.reserve_exact(capacity);
             blob_vec
         }
     }
@@ -55,14 +55,14 @@ impl BlobVec {
         self.capacity
     }
 
-    pub fn reserve(&mut self, amount: usize) {
+    pub fn reserve_exact(&mut self, additional: usize) {
         let available_space = self.capacity - self.len;
-        if available_space < amount {
-            self.grow(amount - available_space);
+        if available_space < additional {
+            self.grow_exact(additional - available_space);
         }
     }
 
-    fn grow(&mut self, increment: usize) {
+    fn grow_exact(&mut self, increment: usize) {
         debug_assert!(self.item_layout.size() != 0);
 
         let new_capacity = self.capacity + increment;
@@ -102,7 +102,7 @@ impl BlobVec {
     /// the newly allocated space must be immediately populated with a valid value
     #[inline]
     pub unsafe fn push_uninit(&mut self) -> usize {
-        self.reserve(1);
+        self.reserve_exact(1);
         let index = self.len;
         self.len += 1;
         index

crates/bevy_ecs/src/storage/table.rs

@@ -1,5 +1,4 @@
 use crate::{
-    archetype::ArchetypeId,
     component::{ComponentId, ComponentInfo, ComponentTicks, Components},
     entity::Entity,
     storage::{BlobVec, SparseSet},
@@ -48,12 +47,24 @@ impl Column {
         }
     }
 
+    #[inline]
+    fn ticks_mut(&mut self) -> &mut Vec<ComponentTicks> {
+        self.ticks.get_mut()
+    }
+
+    /// # Safety
+    /// Assumes data has already been allocated for the given row.
+    #[inline]
+    pub unsafe fn set_unchecked(&mut self, row: usize, data: *mut u8, ticks: ComponentTicks) {
+        self.set_data_unchecked(row, data);
+        *self.ticks_mut().get_unchecked_mut(row) = ticks;
+    }
+
     /// # Safety
-    /// Assumes data has already been allocated for the given row/column.
-    /// Allows aliased mutable accesses to the data at the given `row`. Caller must ensure that this
-    /// does not happen.
+    /// Assumes data has already been allocated for the given row.
     #[inline]
-    pub unsafe fn set_unchecked(&self, row: usize, data: *mut u8) {
+    pub unsafe fn set_data_unchecked(&mut self, row: usize, data: *mut u8) {
+        debug_assert!(row < self.len());
         self.data.set_unchecked(row, data);
     }
 
@@ -68,20 +79,19 @@ impl Column {
     }
 
     /// # Safety
-    /// Assumes data has already been allocated for the given row/column.
-    /// Allows aliased mutable accesses to the row's [ComponentTicks].
-    /// Caller must ensure that this does not happen.
+    /// Assumes data has already been allocated for the given row.
     #[inline]
-    #[allow(clippy::mut_from_ref)]
-    pub unsafe fn get_ticks_unchecked_mut(&self, row: usize) -> &mut ComponentTicks {
+    pub unsafe fn get_ticks_unchecked_mut(&mut self, row: usize) -> &mut ComponentTicks {
         debug_assert!(row < self.len());
-        (*self.ticks.get()).get_unchecked_mut(row)
+        self.ticks_mut().get_unchecked_mut(row)
     }
 
+    /// # Safety
+    /// index must be in-bounds
     #[inline]
     pub(crate) unsafe fn swap_remove_unchecked(&mut self, row: usize) {
         self.data.swap_remove_and_drop_unchecked(row);
-        (*self.ticks.get()).swap_remove(row);
+        self.ticks_mut().swap_remove(row);
     }
 
     #[inline]
@@ -90,26 +100,22 @@ impl Column {
         row: usize,
     ) -> (*mut u8, ComponentTicks) {
         let data = self.data.swap_remove_and_forget_unchecked(row);
-        let ticks = (*self.ticks.get()).swap_remove(row);
+        let ticks = self.ticks_mut().swap_remove(row);
         (data, ticks)
     }
 
-    /// # Safety
-    /// allocated value must be immediately set at the returned row
-    pub(crate) unsafe fn push_uninit(&mut self) -> usize {
+    // # Safety
+    // - ptr must point to valid data of this column's component type
+    pub(crate) unsafe fn push(&mut self, ptr: *mut u8, ticks: ComponentTicks) {
         let row = self.data.push_uninit();
-        (*self.ticks.get()).push(ComponentTicks::new(0));
-        row
+        self.data.set_unchecked(row, ptr);
+        self.ticks_mut().push(ticks);
     }
 
     #[inline]
-    pub(crate) fn reserve(&mut self, additional: usize) {
-        self.data.reserve(additional);
-        // SAFE: unique access to self
-        unsafe {
-            let ticks = &mut (*self.ticks.get());
-            ticks.reserve(additional);
-        }
+    pub(crate) fn reserve_exact(&mut self, additional: usize) {
+        self.data.reserve_exact(additional);
+        self.ticks_mut().reserve_exact(additional);
     }
 
     /// # Safety
@@ -144,8 +150,7 @@ impl Column {
 
     #[inline]
     pub(crate) fn check_change_ticks(&mut self, change_tick: u32) {
-        let ticks = unsafe { (*self.ticks.get()).iter_mut() };
-        for component_ticks in ticks {
+        for component_ticks in self.ticks_mut() {
             component_ticks.check_ticks(change_tick);
         }
     }
@@ -154,29 +159,20 @@ impl Column {
 pub struct Table {
     columns: SparseSet<ComponentId, Column>,
     entities: Vec<Entity>,
-    archetypes: Vec<ArchetypeId>,
-    grow_amount: usize,
-    capacity: usize,
 }
 
 impl Table {
-    pub const fn new(grow_amount: usize) -> Table {
+    pub const fn new() -> Table {
         Self {
             columns: SparseSet::new(),
             entities: Vec::new(),
-            archetypes: Vec::new(),
-            grow_amount,
-            capacity: 0,
         }
     }
 
-    pub fn with_capacity(capacity: usize, column_capacity: usize, grow_amount: usize) -> Table {
+    pub fn with_capacity(capacity: usize, column_capacity: usize) -> Table {
         Self {
             columns: SparseSet::with_capacity(column_capacity),
             entities: Vec::with_capacity(capacity),
-            archetypes: Vec::new(),
-            grow_amount,
-            capacity,
         }
     }
 
@@ -185,14 +181,10 @@ impl Table {
         &self.entities
     }
 
-    pub fn add_archetype(&mut self, archetype_id: ArchetypeId) {
-        self.archetypes.push(archetype_id);
-    }
-
     pub fn add_column(&mut self, component_info: &ComponentInfo) {
         self.columns.insert(
             component_info.id(),
-            Column::with_capacity(component_info, self.capacity()),
+            Column::with_capacity(component_info, self.entities.capacity()),
         )
     }
 
@@ -232,8 +224,7 @@ impl Table {
         for column in self.columns.values_mut() {
             let (data, ticks) = column.swap_remove_and_forget_unchecked(row);
             if let Some(new_column) = new_table.get_column_mut(column.component_id) {
-                new_column.set_unchecked(new_row, data);
-                *new_column.get_ticks_unchecked_mut(new_row) = ticks;
+                new_column.set_unchecked(new_row, data, ticks);
             }
         }
         TableMoveResult {
@@ -263,8 +254,7 @@ impl Table {
         for column in self.columns.values_mut() {
             if let Some(new_column) = new_table.get_column_mut(column.component_id) {
                 let (data, ticks) = column.swap_remove_and_forget_unchecked(row);
-                new_column.set_unchecked(new_row, data);
-                *new_column.get_ticks_unchecked_mut(new_row) = ticks;
+                new_column.set_unchecked(new_row, data, ticks);
             } else {
                 column.swap_remove_unchecked(row);
             }
@@ -296,8 +286,7 @@ impl Table {
         for column in self.columns.values_mut() {
             let new_column = new_table.get_column_mut(column.component_id).unwrap();
             let (data, ticks) = column.swap_remove_and_forget_unchecked(row);
-            new_column.set_unchecked(new_row, data);
-            *new_column.get_ticks_unchecked_mut(new_row) = ticks;
+            new_column.set_unchecked(new_row, data, ticks);
         }
         TableMoveResult {
             new_row,
@@ -324,20 +313,16 @@ impl Table {
         self.columns.contains(component_id)
     }
 
-    pub fn reserve(&mut self, amount: usize) {
-        let available_space = self.capacity - self.len();
-        if available_space < amount {
-            let min_capacity = self.len() + amount;
-            // normally we would check if min_capacity is 0 for the below calculation, but amount >
-            // available_space and available_space > 0, so min_capacity > 1
-            let new_capacity =
-                ((min_capacity + self.grow_amount - 1) / self.grow_amount) * self.grow_amount;
-            let reserve_amount = new_capacity - self.len();
+    pub fn reserve(&mut self, additional: usize) {
+        if self.entities.capacity() - self.entities.len() < additional {
+            self.entities.reserve(additional);
+
+            // use entities vector capacity as driving capacity for all related allocations
+            let new_capacity = self.entities.capacity();
+
             for column in self.columns.values_mut() {
-                column.reserve(reserve_amount);
+                column.reserve_exact(new_capacity - column.len());
             }
-            self.entities.reserve(reserve_amount);
-            self.capacity = new_capacity;
         }
     }
 
@@ -358,7 +343,7 @@ impl Table {
 
     #[inline]
     pub fn capacity(&self) -> usize {
-        self.capacity
+        self.entities.capacity()
     }
 
     #[inline]
@@ -389,7 +374,7 @@ pub struct Tables {
 
 impl Default for Tables {
     fn default() -> Self {
-        let empty_table = Table::with_capacity(0, 0, 64);
+        let empty_table = Table::with_capacity(0, 0);
         Tables {
             tables: vec![empty_table],
             table_ids: HashMap::default(),
@@ -446,7 +431,7 @@ impl Tables {
         let hash = hasher.finish();
         let tables = &mut self.tables;
         *self.table_ids.entry(hash).or_insert_with(move || {
-            let mut table = Table::with_capacity(0, component_ids.len(), 64);
+            let mut table = Table::with_capacity(0, component_ids.len());
             for component_id in component_ids.iter() {
                 table.add_column(components.get_info_unchecked(*component_id));
             }
@@ -496,18 +481,19 @@ mod tests {
         let type_info = TypeInfo::of::<usize>();
         let component_id = components.get_or_insert_with(type_info.type_id(), || type_info);
         let columns = &[component_id];
-        let mut table = Table::with_capacity(0, columns.len(), 64);
+        let mut table = Table::with_capacity(0, columns.len());
         table.add_column(components.get_info(component_id).unwrap());
         let entities = (0..200).map(Entity::new).collect::<Vec<_>>();
-        for (row, entity) in entities.iter().cloned().enumerate() {
+        for entity in entities.iter() {
+            // SAFE: we allocate and immediately set data afterwards
             unsafe {
-                table.allocate(entity);
+                let row = table.allocate(*entity);
                 let mut value = row;
                 let value_ptr = ((&mut value) as *mut usize).cast::<u8>();
                 table
-                    .get_column(component_id)
+                    .get_column_mut(component_id)
                     .unwrap()
-                    .set_unchecked(row, value_ptr);
+                    .set_data_unchecked(row, value_ptr);
             };
         }
 

crates/bevy_ecs/src/world/entity_ref.rs

@@ -275,7 +275,7 @@ impl<'w> EntityMut<'w> {
         };
         self.location = new_location;
 
-        let table = &storages.tables[archetype.table_id()];
+        let table = &mut storages.tables[archetype.table_id()];
         let table_row = archetype.entity_table_row(new_location.index);
         // SAFE: table row is valid
         unsafe {

crates/bevy_ecs/src/world/mod.rs

@@ -718,12 +718,10 @@ impl World {
         let column = unique_components
             .get_mut(component_id)
             .unwrap_or_else(|| panic!("resource does not exist: {}", std::any::type_name::<T>()));
-        // SAFE: new location is immediately written to below
-        let row = unsafe { column.push_uninit() };
-        // SAFE: row was just allocated above
-        unsafe { column.set_unchecked(row, ptr) };
-        // SAFE: row was just allocated above
-        unsafe { *column.get_ticks_unchecked_mut(row) = ticks };
+        unsafe {
+            // SAFE: pointer is of type T
+            column.push(ptr, ticks);
+        }
         result
     }
 
@@ -787,13 +785,8 @@ impl World {
         if column.is_empty() {
             // SAFE: column is of type T and has been allocated above
             let data = (&mut value as *mut T).cast::<u8>();
-            // SAFE: new location is immediately written to below
-            let row = column.push_uninit();
-            // SAFE: index was just allocated above
-            column.set_unchecked(row, data);
             std::mem::forget(value);
-            // SAFE: index was just allocated above
-            *column.get_ticks_unchecked_mut(row) = ComponentTicks::new(change_tick);
+            column.push(data, ComponentTicks::new(change_tick));
         } else {
             // SAFE: column is of type T and has already been allocated
             *column.get_unchecked(0).cast::<T>() = value;