Remove redundent information and optimize dynamic allocations in Table

crates/bevy_ecs/src/storage/blob_array.rs

@@ -333,17 +333,41 @@ impl BlobArray {
         index_to_keep: usize,
     ) -> OwningPtr<'_> {
         if index_to_remove != index_to_keep {
-            std::ptr::swap_nonoverlapping::<u8>(
-                self.get_unchecked_mut(index_to_keep).as_ptr(),
-                self.get_unchecked_mut(index_to_remove).as_ptr(),
-                self.item_layout.size(),
-            );
+            return self
+                .swap_remove_and_forget_unchecked_nonoverlapping(index_to_remove, index_to_keep);
         }
         // Now the element that used to be in index `index_to_remove` is now in index `index_to_keep` (after swap)
         // If we are storing ZSTs than the index doesn't actually matter because the size is 0.
         self.get_unchecked_mut(index_to_keep).promote()
     }
 
+    /// The same as [`Self::swap_remove_and_forget_unchecked`] but the two elements must non-overlapping.
+    ///
+    /// # Safety
+    /// The caller must ensure that:
+    /// - `index_to_keep` < `len`
+    /// - `index_to_remove` < `len`
+    /// - `index_to_remove` != `index_to_keep`
+    /// - If `index_to_keep` == `len` - 1, and the caller has the length saved, update the length to reflect that the element with index
+    /// `len` - 1 is not valid to use (set `len` to `len` - 1).
+    /// - If the length wasn't updated by the caller, they must use [`Self::initialize_unchecked`] to initialize an element in the index `index_to_keep`,
+    /// because after calling this method, the element with index `index_to_keep` will not be valid to use.
+    #[inline]
+    pub unsafe fn swap_remove_and_forget_unchecked_nonoverlapping(
+        &mut self,
+        index_to_remove: usize,
+        index_to_keep: usize,
+    ) -> OwningPtr<'_> {
+        std::ptr::swap_nonoverlapping::<u8>(
+            self.get_unchecked_mut(index_to_keep).as_ptr(),
+            self.get_unchecked_mut(index_to_remove).as_ptr(),
+            self.item_layout.size(),
+        );
+        // Now the element that used to be in index `index_to_remove` is now in index `index_to_keep` (after swap)
+        // If we are storing ZSTs than the index doesn't actually matter because the size is 0.
+        self.get_unchecked_mut(index_to_keep).promote()
+    }
+
     /// This method will swap two elements in the array, and drop the one with the index `index_to_remove`.
     ///
     /// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.
@@ -373,4 +397,29 @@ impl BlobArray {
             drop(value);
         }
     }
+
+    /// The same as [`Self::swap_remove_and_drop_unchecked`] but the two elements must non-overlapping.
+    ///
+    /// # Safety
+    /// The caller must ensure that:
+    /// - `index_to_keep` < `len`
+    /// - `index_to_remove` < `len`
+    /// - `index_to_remove` != `index_to_keep`
+    /// - If `index_to_keep` == `len` - 1, and the caller has the length saved, update the length to reflect that the element with index
+    /// `len` - 1 is not valid to use (set `len` to `len` - 1).
+    /// - If the length wasn't updated by the caller, they must use [`Self::initialize_unchecked`] to initialize an element in the index `index_to_keep`,
+    /// because after calling this method, the element with index `index_to_keep` will not be valid to use.
+    #[inline]
+    pub unsafe fn swap_remove_and_drop_unchecked_nonoverlapping(
+        &mut self,
+        index_to_remove: usize,
+        index_to_keep: usize,
+    ) {
+        let drop = self.drop;
+        let value =
+            self.swap_remove_and_forget_unchecked_nonoverlapping(index_to_remove, index_to_keep);
+        if let Some(drop) = drop {
+            drop(value);
+        }
+    }
 }

crates/bevy_ecs/src/storage/table/column.rs

@@ -26,6 +26,30 @@ impl ThinColumn {
         }
     }
 
+    /// Swap-remove and drop the removed element, but the component at `row` must not be the last elment.
+    ///
+    /// # Safety
+    /// - `row.as_usize()` < `len`
+    /// - `last_element_index` = `len - 1`
+    /// - `last_element_index` != `row.as_usize()`
+    /// - _update the `len`_ to `len - 1`, or immediately initialize another element in the `last_element_index`
+    pub unsafe fn swap_remove_and_drop_unchecked_nonoverlapping(
+        &mut self,
+        last_element_index: usize,
+        row: TableRow,
+    ) {
+        self.data
+            .swap_remove_and_drop_unchecked_nonoverlapping(row.as_usize(), last_element_index);
+        let added = &mut self
+            .added_ticks
+            .swap_remove_and_forget_unchecked_nonoverlapping(row.as_usize(), last_element_index);
+        let changed = &mut self
+            .changed_ticks
+            .swap_remove_and_forget_unchecked_nonoverlapping(row.as_usize(), last_element_index);
+        std::ptr::drop_in_place(added as *mut UnsafeCell<Tick>);
+        std::ptr::drop_in_place(changed as *mut UnsafeCell<Tick>);
+    }
+
     /// Swap-remove and drop the removed element.
     ///
     /// # Safety

crates/bevy_ecs/src/storage/table/mod.rs

@@ -151,7 +151,7 @@ pub(crate) struct TableBuilder {
 }
 
 impl TableBuilder {
-    ///  Start building a new [`Table`] with a specified `column_capacity` (How many components per column?) and a `capacity` (How many columns?)
+    /// Start building a new [`Table`] with a specified `column_capacity` (How many components per column?) and a `capacity` (How many columns?)
     pub fn with_capacity(capacity: usize, column_capacity: usize) -> Self {
         Self {
             columns: SparseSet::with_capacity(column_capacity),
@@ -175,7 +175,6 @@ impl TableBuilder {
         Table {
             columns: self.columns.into_immutable(),
             entities: Vec::with_capacity(self.capacity),
-            poisoned: false,
         }
     }
 }
@@ -195,8 +194,15 @@ impl TableBuilder {
 pub struct Table {
     columns: ImmutableSparseSet<ComponentId, ThinColumn>,
     entities: Vec<Entity>,
-    /// This is set to `true` if an allocation triggered an unwind
-    poisoned: bool,
+}
+
+struct AbortOnPanic;
+
+impl Drop for AbortOnPanic {
+    fn drop(&mut self) {
+        // Panicking while unwinding will force an abort.
+        panic!("Aborting due to allocator error");
+    }
 }
 
 impl Table {
@@ -227,12 +233,17 @@ impl Table {
     pub(crate) unsafe fn swap_remove_unchecked(&mut self, row: TableRow) -> Option<Entity> {
         debug_assert!(row.as_usize() < self.len());
         let last_element_index = self.len() - 1;
-        for col in self.columns.values_mut() {
-            // SAFETY:
-            // - `row` < `len`
-            // - `last_element_index` = `len` - 1
-            // - the `len` is kept within `self.entities`, it will update accordingly.
-            unsafe { col.swap_remove_and_drop_unchecked(last_element_index, row) };
+        if row.as_usize() != last_element_index {
+            for col in self.columns.values_mut() {
+                // SAFETY:
+                // - `row` < `len`
+                // - `last_element_index` = `len` - 1
+                // - `row` != last_elment_index
+                // - the `len` is kept within `self.entities`, it will update accordingly.
+                unsafe {
+                    col.swap_remove_and_drop_unchecked_nonoverlapping(last_element_index, row)
+                };
+            }
         }
         let is_last = row.as_usize() == last_element_index;
         self.entities.swap_remove(row.as_usize());
@@ -268,7 +279,6 @@ impl Table {
                 column.swap_remove_and_forget_unchecked(last_element_index, row);
             }
         }
-
         TableMoveResult {
             new_row,
             swapped_entity: if is_last {
@@ -529,14 +539,14 @@ impl Table {
     ///
     /// The current capacity of the columns should be 0, if it's not 0, then the previous data will be overwritten and leaked.
     fn alloc_columns(&mut self, new_capacity: NonZeroUsize) {
-        self.poisoned = true;
         // If any of these allocations trigger an unwind, the wrong capacity will be used while dropping this table - UB.
-        // To avoid this, we manually set `self.poisoned` to true. If the drop method is called when `self.poisoned` is true,
-        // nothing will be deallocated - avoiding possible UB (but still leaking memory).
+        // To avoid this, we use `AbortOnPanic`. If the allocation triggered a panic, the `AbortOnPanic`'s Drop impl will be
+        // called, and abort the program.
+        let _guard = AbortOnPanic;
         for col in self.columns.values_mut() {
             col.alloc(new_capacity);
         }
-        self.poisoned = false; // The allocations didn't trigger an unwind, so we set the flag to `false`.
+        core::mem::forget(_guard); // The allocation was successful, so we don't drop the guard.
     }
 
     /// Reallocate memory for the columns in the [`Table`]
@@ -548,10 +558,10 @@ impl Table {
         current_column_capacity: NonZeroUsize,
         new_capacity: NonZeroUsize,
     ) {
-        self.poisoned = true;
         // If any of these allocations trigger an unwind, the wrong capacity will be used while dropping this table - UB.
-        // To avoid this, we manually set `self.poisoned` to true. If the drop method is called when `self.poisoned` is true,
-        // nothing will be deallocated - avoiding possible UB (but still leaking memory).
+        // To avoid this, we use `AbortOnPanic`. If the allocation triggered a panic, the `AbortOnPanic`'s Drop impl will be
+        // called, and abort the program.
+        let _guard = AbortOnPanic;
 
         // SAFETY:
         // - There's no overflow
@@ -560,7 +570,7 @@ impl Table {
         for col in self.columns.values_mut() {
             col.realloc(current_column_capacity, new_capacity);
         }
-        self.poisoned = false; // The allocations didn't trigger an unwind, so we set the flag to `false`.
+        core::mem::forget(_guard); // The allocation was successful, so we don't drop the guard.
     }
 
     /// Allocates space for a new entity
@@ -791,15 +801,13 @@ impl IndexMut<TableId> for Tables {
 
 impl Drop for Table {
     fn drop(&mut self) {
-        if !self.poisoned {
-            let len = self.len();
-            let cap = self.capacity();
-            self.entities.clear();
-            for col in self.columns.values_mut() {
-                // SAFETY: `cap` and `len` are correct
-                unsafe {
-                    col.drop(cap, len);
-                }
+        let len = self.len();
+        let cap = self.capacity();
+        self.entities.clear();
+        for col in self.columns.values_mut() {
+            // SAFETY: `cap` and `len` are correct
+            unsafe {
+                col.drop(cap, len);
             }
         }
     }