Use ThinColumn in ComponentSparseSet

crates/bevy_ecs/src/storage/blob_vec.rs

@@ -1,7 +1,7 @@
 use alloc::alloc::handle_alloc_error;
 use bevy_ptr::{OwningPtr, Ptr, PtrMut};
 use bevy_utils::OnDrop;
-use core::{alloc::Layout, cell::UnsafeCell, num::NonZero, ptr::NonNull};
+use core::{alloc::Layout, num::NonZero, ptr::NonNull};
 
 /// A flat, type-erased data storage type
 ///
@@ -88,12 +88,6 @@ impl BlobVec {
         self.len == 0
     }
 
-    /// Returns the [`Layout`] of the element type stored in the vector.
-    #[inline]
-    pub fn layout(&self) -> Layout {
-        self.item_layout
-    }
-
     /// Reserves the minimum capacity for at least `additional` more elements to be inserted in the given `BlobVec`.
     /// After calling `reserve_exact`, capacity will be greater than or equal to `self.len() + additional`. Does nothing if
     /// the capacity is already sufficient.
@@ -251,7 +245,7 @@ impl BlobVec {
     /// Appends an element to the back of the vector.
     ///
     /// # Safety
-    /// The `value` must match the [`layout`](`BlobVec::layout`) of the elements in the [`BlobVec`].
+    /// The `value` must match the layout of the elements in the [`BlobVec`].
     #[inline]
     pub unsafe fn push(&mut self, value: OwningPtr<'_>) {
         self.reserve(1);
@@ -293,22 +287,6 @@ impl BlobVec {
         unsafe { p.promote() }
     }
 
-    /// Removes the value at `index` and drops it.
-    /// Does not do any bounds checking on `index`.
-    /// The removed element is replaced by the last element of the `BlobVec`.
-    ///
-    /// # Safety
-    /// It is the caller's responsibility to ensure that `index` is `< self.len()`.
-    #[inline]
-    pub unsafe fn swap_remove_and_drop_unchecked(&mut self, index: usize) {
-        debug_assert!(index < self.len());
-        let drop = self.drop;
-        let value = self.swap_remove_and_forget_unchecked(index);
-        if let Some(drop) = drop {
-            drop(value);
-        }
-    }
-
     /// Returns a reference to the element at `index`, without doing bounds checking.
     ///
     /// # Safety
@@ -357,22 +335,6 @@ impl BlobVec {
         unsafe { PtrMut::new(self.data) }
     }
 
-    /// Get a reference to the entire [`BlobVec`] as if it were an array with elements of type `T`
-    ///
-    /// # Safety
-    /// The type `T` must be the type of the items in this [`BlobVec`].
-    pub unsafe fn get_slice<T>(&self) -> &[UnsafeCell<T>] {
-        // SAFETY: the inner data will remain valid for as long as 'self.
-        unsafe { core::slice::from_raw_parts(self.data.as_ptr() as *const UnsafeCell<T>, self.len) }
-    }
-
-    /// Returns the drop function for values stored in the vector,
-    /// or `None` if they don't need to be dropped.
-    #[inline]
-    pub fn get_drop(&self) -> Option<unsafe fn(OwningPtr<'_>)> {
-        self.drop
-    }
-
     /// Clears the vector, removing (and dropping) all values.
     ///
     /// Note that this method has no effect on the allocated capacity of the vector.

crates/bevy_ecs/src/storage/mod.rs

@@ -60,3 +60,24 @@ impl Storages {
         }
     }
 }
+
+/// Calls a function, aborting the program if it panics.
+/// This should be used whenever an operation may cause a panic
+/// while an object is still in an invalid state, which could cause UB when dropped.
+#[inline(always)]
+fn abort_on_panic<F: FnOnce() -> R, R>(f: F) -> R {
+    struct AbortOnPanic;
+
+    impl Drop for AbortOnPanic {
+        fn drop(&mut self) {
+            // Panicking while unwinding will force an abort.
+            panic!("Aborting due to allocator error");
+        }
+    }
+
+    let _guard = AbortOnPanic;
+    let ret = f();
+    // the operation was successful, so we don't drop the guard
+    core::mem::forget(_guard);
+    ret
+}

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -2,13 +2,15 @@ use crate::{
     change_detection::MaybeLocation,
     component::{ComponentId, ComponentInfo, ComponentTicks, Tick, TickCells},
     entity::Entity,
-    storage::{Column, TableRow},
+    storage::TableRow,
 };
 use alloc::{boxed::Box, vec::Vec};
 use bevy_ptr::{OwningPtr, Ptr};
-use core::{cell::UnsafeCell, hash::Hash, marker::PhantomData, panic::Location};
+use core::{cell::UnsafeCell, hash::Hash, marker::PhantomData, num::NonZeroUsize, panic::Location};
 use nonmax::NonMaxUsize;
 
+use super::{abort_on_panic, ThinColumn};
+
 type EntityIndex = u32;
 
 #[derive(Debug)]
@@ -116,7 +118,8 @@ impl<I: SparseSetIndex, V> SparseArray<I, V> {
 /// Designed for relatively fast insertions and deletions.
 #[derive(Debug)]
 pub struct ComponentSparseSet {
-    dense: Column,
+    /// SAFETY: Equal in length & capacity to `self.entities`
+    dense: ThinColumn,
     // Internally this only relies on the Entity index to keep track of where the component data is
     // stored for entities that are alive. The generation is not required, but is stored
     // in debug builds to validate that access is correct.
@@ -132,29 +135,37 @@ impl ComponentSparseSet {
     /// initial `capacity`.
     pub(crate) fn new(component_info: &ComponentInfo, capacity: usize) -> Self {
         Self {
-            dense: Column::with_capacity(component_info, capacity),
+            dense: ThinColumn::with_capacity(component_info, capacity),
             entities: Vec::with_capacity(capacity),
             sparse: Default::default(),
         }
     }
 
     /// Removes all of the values stored within.
     pub(crate) fn clear(&mut self) {
-        self.dense.clear();
+        // SAFETY: `self.len()` is the length of the column
+        unsafe {
+            self.dense.clear(self.len());
+        }
         self.entities.clear();
         self.sparse.clear();
     }
 
     /// Returns the number of component values in the sparse set.
     #[inline]
     pub fn len(&self) -> usize {
-        self.dense.len()
+        self.entities.len()
     }
 
     /// Returns `true` if the sparse set contains no component values.
     #[inline]
     pub fn is_empty(&self) -> bool {
-        self.dense.len() == 0
+        self.len() == 0
+    }
+
+    #[inline]
+    fn capacity(&self) -> usize {
+        self.entities.capacity()
     }
 
     /// Inserts the `entity` key and component `value` pair into this sparse
@@ -175,13 +186,16 @@ impl ComponentSparseSet {
             assert_eq!(entity, self.entities[dense_index.as_usize()]);
             self.dense.replace(dense_index, value, change_tick, caller);
         } else {
-            let dense_index = self.dense.len();
-            self.dense
-                .push(value, ComponentTicks::new(change_tick), caller);
-            self.sparse
-                .insert(entity.index(), TableRow::from_usize(dense_index));
+            let dense_index = TableRow::from_usize(self.len());
+            self.reserve(1);
+            // SAFETY: `dense_index` is the last element in the column after the call to `reserve`
+            unsafe {
+                self.dense
+                    .initialize(dense_index, value, change_tick, caller);
+            }
+            self.sparse.insert(entity.index(), dense_index);
             #[cfg(debug_assertions)]
-            assert_eq!(self.entities.len(), dense_index);
+            assert_eq!(self.len(), dense_index.as_usize());
             #[cfg(not(debug_assertions))]
             self.entities.push(entity.index());
             #[cfg(debug_assertions)]
@@ -314,10 +328,14 @@ impl ComponentSparseSet {
         self.sparse.remove(entity.index()).map(|dense_index| {
             #[cfg(debug_assertions)]
             assert_eq!(entity, self.entities[dense_index.as_usize()]);
+            let last_index = self.len() - 1;
+            let is_last = dense_index.as_usize() == last_index;
             self.entities.swap_remove(dense_index.as_usize());
-            let is_last = dense_index.as_usize() == self.dense.len() - 1;
             // SAFETY: dense_index was just removed from `sparse`, which ensures that it is valid
-            let (value, _, _) = unsafe { self.dense.swap_remove_and_forget_unchecked(dense_index) };
+            let value = unsafe {
+                self.dense
+                    .swap_remove_and_forget_unchecked(last_index, dense_index)
+            };
             if !is_last {
                 let swapped_entity = self.entities[dense_index.as_usize()];
                 #[cfg(not(debug_assertions))]
@@ -337,11 +355,13 @@ impl ComponentSparseSet {
         if let Some(dense_index) = self.sparse.remove(entity.index()) {
             #[cfg(debug_assertions)]
             assert_eq!(entity, self.entities[dense_index.as_usize()]);
+            let last_index = self.len() - 1;
+            let is_last = dense_index.as_usize() == last_index;
             self.entities.swap_remove(dense_index.as_usize());
-            let is_last = dense_index.as_usize() == self.dense.len() - 1;
             // SAFETY: if the sparse index points to something in the dense vec, it exists
             unsafe {
-                self.dense.swap_remove_unchecked(dense_index);
+                self.dense
+                    .swap_remove_and_drop_unchecked(last_index, dense_index);
             }
             if !is_last {
                 let swapped_entity = self.entities[dense_index.as_usize()];
@@ -358,7 +378,56 @@ impl ComponentSparseSet {
     }
 
     pub(crate) fn check_change_ticks(&mut self, change_tick: Tick) {
-        self.dense.check_change_ticks(change_tick);
+        // SAFETY: `self.len()` is equal to the length of the column
+        unsafe {
+            self.dense.check_change_ticks(self.len(), change_tick);
+        }
+    }
+
+    /// Reserves `additional` elements worth of capacity within the table.
+    pub(crate) fn reserve(&mut self, additional: usize) {
+        if self.capacity() - self.len() < additional {
+            let column_cap = self.capacity();
+            self.entities.reserve(additional);
+
+            // use entities vector capacity as driving capacity for all related allocations
+            //
+            // SAFETY: `additional` must be > 0 for the branch above to succeed, so `new_capacity`
+            // must be nonzero
+            let new_capacity = unsafe { NonZeroUsize::new_unchecked(self.capacity()) };
+
+            match NonZeroUsize::new(column_cap) {
+                Some(column_cap) => {
+                    // If any of these allocations trigger an unwind, the wrong capacity will be used while dropping this table - UB.
+                    // To avoid this, we use `abort_on_panic`. If the allocation triggered a panic, the guard will be triggered, and
+                    // abort the program.
+
+                    // SAFETY: `column_cap` is indeed the column's capacity
+                    abort_on_panic(|| unsafe {
+                        self.dense.realloc(column_cap, new_capacity);
+                    });
+                }
+                None => {
+                    // If any of these allocations trigger an unwind, the wrong capacity will be used while dropping this table - UB.
+                    // To avoid this, we use `abort_on_panic`. If the allocation triggered a panic, the guard will be triggered, and
+                    // abort the program.
+                    abort_on_panic(|| {
+                        self.dense.alloc(new_capacity);
+                    });
+                }
+            }
+        }
+    }
+}
+
+impl Drop for ComponentSparseSet {
+    fn drop(&mut self) {
+        // SAFETY:
+        // - self.capacity() is the capacity of the column
+        // - self.len() is the length of the column
+        unsafe {
+            self.dense.drop(self.capacity(), self.len());
+        }
     }
 }