add swap functions to table

Author: janis-bhm

crates/bevy_ecs/src/storage/blob_array.rs

@@ -388,7 +388,7 @@ impl BlobArray {
         self.get_unchecked_mut(index_to_keep).promote()
     }
 
-    /// The same as [`Self::swap_remove_unchecked`] but the two elements must non-overlapping.
+    /// This method will swap two elements in the array.
     ///
     /// # Safety
     /// - `index_to_keep` must be safe to access (within the bounds of the length of the array).
@@ -398,11 +398,11 @@ impl BlobArray {
     ///     1) initialize a different value in `index_to_keep`
     ///     2) update the saved length of the array if `index_to_keep` was the last element.
     #[inline]
-    pub unsafe fn swap_remove_unchecked_nonoverlapping(
+    pub unsafe fn swap_unchecked_nonoverlapping(
         &mut self,
         index_to_remove: usize,
         index_to_keep: usize,
-    ) -> OwningPtr<'_> {
+    ) {
         #[cfg(debug_assertions)]
         {
             debug_assert!(self.capacity > index_to_keep);
@@ -415,6 +415,24 @@ impl BlobArray {
             self.get_unchecked_mut(index_to_remove).as_ptr(),
             self.item_layout.size(),
         );
+    }
+
+    /// The same as [`Self::swap_remove_unchecked`] but the two elements must non-overlapping.
+    ///
+    /// # Safety
+    /// - `index_to_keep` must be safe to access (within the bounds of the length of the array).
+    /// - `index_to_remove` must be safe to access (within the bounds of the length of the array).
+    /// - `index_to_remove` != `index_to_keep`
+    /// -  The caller should address the inconsistent state of the array that has occurred after the swap, either:
+    ///     1) initialize a different value in `index_to_keep`
+    ///     2) update the saved length of the array if `index_to_keep` was the last element.
+    #[inline]
+    pub unsafe fn swap_remove_unchecked_nonoverlapping(
+        &mut self,
+        index_to_remove: usize,
+        index_to_keep: usize,
+    ) -> OwningPtr<'_> {
+        self.swap_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()

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

@@ -36,6 +36,23 @@ impl ThinColumn {
         }
     }
 
+    /// Swap the two elements.
+    ///
+    /// # Safety
+    /// - `a.as_usize()` < `len`
+    /// - `b.as_usize()` < `len`
+    /// - `a` != `b`
+    pub(crate) unsafe fn swap_unchecked(&mut self, a: TableRow, b: TableRow) {
+        let a = a.index();
+        let b = b.index();
+        self.data.swap_unchecked_nonoverlapping(a, b);
+        self.added_ticks.swap_unchecked_nonoverlapping(a, b);
+        self.changed_ticks.swap_unchecked_nonoverlapping(a, b);
+        self.changed_by.as_mut().map(|changed_by| {
+            changed_by.swap_unchecked_nonoverlapping(a, b);
+        });
+    }
+
     /// Swap-remove and drop the removed element, but the component at `row` must not be the last element.
     ///
     /// # Safety

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

@@ -13,7 +13,7 @@ use core::{
     alloc::Layout,
     cell::UnsafeCell,
     num::NonZeroUsize,
-    ops::{Index, IndexMut},
+    ops::{Index, IndexMut, Range},
     panic::Location,
 };
 use nonmax::NonMaxU32;
@@ -167,6 +167,7 @@ impl TableBuilder {
         Table {
             columns: self.columns.into_immutable(),
             entities: self.entities,
+            disabled_entities: 0,
         }
     }
 }
@@ -191,6 +192,7 @@ impl TableBuilder {
 pub struct Table {
     columns: ImmutableSparseSet<ComponentId, ThinColumn>,
     entities: Vec<Entity>,
+    disabled_entities: u32,
 }
 
 struct AbortOnPanic;
@@ -206,7 +208,20 @@ impl Table {
     /// Fetches a read-only slice of the entities stored within the [`Table`].
     #[inline]
     pub fn entities(&self) -> &[Entity] {
-        &self.entities
+        &self.entities[self.disabled_entities as usize..]
+    }
+
+    /// Get the valid table rows (i.e. non-disabled entities).
+    #[inline]
+    pub fn table_rows(&self) -> Range<TableRow> {
+        // SAFETY:
+        // - No entity row may be in more than one table row at once, so there are no duplicates,
+        // and there can not be an entity row of u32::MAX. Therefore, this can not be max either.
+        // - self.disabled_entities <= self.entity_count()
+        unsafe {
+            TableRow::new(NonMaxU32::new_unchecked(self.disabled_entities))
+                ..TableRow::new(NonMaxU32::new_unchecked(self.entity_count()))
+        }
     }
 
     /// Get the capacity of this table, in entities.
@@ -216,6 +231,39 @@ impl Table {
         self.entities.capacity()
     }
 
+    /// Disables the entity at the given row, moving it to the back of the table and returns the entities that were swapped (if any) together with their new `TableRow`.
+    ///
+    /// # Safety
+    /// `row` must be in-bounds (`row.as_usize()` < `self.len()`) and not disabled.
+    pub(crate) unsafe fn swap_disable_unchecked(
+        &mut self,
+        row: TableRow,
+    ) -> Option<((Entity, TableRow), (Entity, TableRow))> {
+        debug_assert!(row.index_u32() < self.entity_count());
+        debug_assert!(row.index_u32() >= self.disabled_entities);
+
+        if row.index_u32() != 0 {
+            // SAFETY: `self.disabled_entities` is always less than `u32::MAX`,
+            // as guaranteed by `allocate`.
+            let other = TableRow::new(unsafe { NonMaxU32::new_unchecked(self.disabled_entities) });
+
+            for col in self.columns.values_mut() {
+                col.swap_unchecked(row, other);
+            }
+
+            self.entities.swap(row.index(), other.index());
+            self.disabled_entities += 1;
+
+            Some((
+                (*self.entities.get_unchecked(other.index()), other),
+                (*self.entities.get_unchecked(row.index()), row),
+            ))
+        } else {
+            self.disabled_entities += 1;
+            None
+        }
+    }
+
     /// Removes the entity at the given row and returns the entity swapped in to replace it (if an
     /// entity was swapped in)
     ///

crates/bevy_ecs/src/storage/thin_array_ptr.rs

@@ -181,6 +181,24 @@ impl<T> ThinArrayPtr<T> {
         }
     }
 
+    /// Perform a [`swap`](https://doc.rust-lang.org/std/primitive.slice.html#method.swap).
+    ///
+    /// # Safety
+    /// - `a` must be safe to access (within the bounds of the length of the array).
+    /// - `b` must be safe to access (within the bounds of the length of the array).
+    /// - `a` != `b`
+    #[inline]
+    pub unsafe fn swap_unchecked_nonoverlapping(&mut self, a: usize, b: usize) {
+        #[cfg(debug_assertions)]
+        {
+            debug_assert!(self.capacity > a);
+            debug_assert!(self.capacity > b);
+            debug_assert_ne!(a, b);
+        }
+        let base_ptr = self.data.as_ptr();
+        ptr::copy_nonoverlapping(base_ptr.add(a), base_ptr.add(b), 1);
+    }
+
     /// Perform a [`swap-remove`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.swap_remove) and return the removed value.
     ///
     /// # Safety