disabling

Author: janis-bhm

crates/bevy_ecs/src/archetype.rs

@@ -327,6 +327,7 @@ impl Edges {
 }
 
 /// Metadata about an [`Entity`] in a [`Archetype`].
+#[derive(Clone, Copy)]
 pub struct ArchetypeEntity {
     entity: Entity,
     table_row: TableRow,
@@ -634,6 +635,52 @@ impl Archetype {
         self.entities.reserve(additional);
     }
 
+    /// Disables the entity at `row` by swapping it with the first enabled
+    /// entity. Returns the swapped entities with their respective table and
+    /// archetype rows, or `None` if no swap occurred.
+    ///
+    /// # Panics
+    /// This function will panic if `row >= self.entities.len()`
+    #[inline]
+    pub(crate) fn swap_disable(
+        &mut self,
+        row: ArchetypeRow,
+    ) -> (
+        (ArchetypeEntity, ArchetypeRow),
+        Option<(ArchetypeEntity, ArchetypeRow)>,
+    ) {
+        debug_assert!(row.index_u32() < self.len());
+        debug_assert!(row.index_u32() >= self.disabled_entities);
+
+        if row.index_u32() == self.disabled_entities {
+            // no need to swap, just increment the disabled count
+            self.disabled_entities += 1;
+
+            // SAFETY: `row` is guaranteed to be in-bounds.
+            (
+                (unsafe { *self.entities.get_unchecked(row.index()) }, row),
+                None,
+            )
+        } else {
+            // SAFETY: `self.disabled_entities` is always less than `u32::MAX`, as guaranteed by `allocate`.
+            let disabled_row =
+                ArchetypeRow::new(unsafe { NonMaxU32::new_unchecked(self.disabled_entities) });
+
+            self.entities.swap(row.index(), disabled_row.index());
+
+            // SAFETY: Both `row` and `other` are guaranteed to be in-bounds.
+            unsafe {
+                (
+                    (
+                        *self.entities.get_unchecked(disabled_row.index()),
+                        disabled_row,
+                    ),
+                    Some((*self.entities.get_unchecked(row.index()), row)),
+                )
+            }
+        }
+    }
+
     /// Removes the entity at `row` by swapping it out. Returns the table row the entity is stored
     /// in.
     ///

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

@@ -224,29 +224,33 @@ impl Table {
     pub(crate) unsafe fn swap_disable_unchecked(
         &mut self,
         row: TableRow,
-    ) -> Option<((Entity, TableRow), (Entity, TableRow))> {
+    ) -> ((Entity, TableRow), Option<(Entity, TableRow)>) {
         debug_assert!(row.index_u32() < self.len());
         debug_assert!(row.index_u32() >= self.disabled_entities);
 
-        if row.index_u32() != 0 {
+        if row.index_u32() != self.disabled_entities {
             // 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) });
+            let disabled_row =
+                TableRow::new(unsafe { NonMaxU32::new_unchecked(self.disabled_entities) });
 
             for col in self.columns.values_mut() {
-                col.swap_unchecked(row, other);
+                col.swap_unchecked(row, disabled_row);
             }
 
-            self.entities.swap(row.index(), other.index());
+            self.entities.swap(row.index(), disabled_row.index());
             self.disabled_entities += 1;
 
-            Some((
-                (*self.entities.get_unchecked(other.index()), other),
-                (*self.entities.get_unchecked(row.index()), row),
-            ))
+            (
+                (
+                    *self.entities.get_unchecked(disabled_row.index()),
+                    disabled_row,
+                ),
+                Some((*self.entities.get_unchecked(row.index()), row)),
+            )
         } else {
             self.disabled_entities += 1;
-            None
+            ((*self.entities.get_unchecked(row.index()), row), None)
         }
     }
 

crates/bevy_ecs/src/world/entity_ref.rs

@@ -2562,6 +2562,76 @@ impl<'w> EntityWorldMut<'w> {
         self.despawn_with_caller(MaybeLocation::caller());
     }
 
+    /// Disable the current entity.
+    ///
+    pub fn disable(self) -> EntityLocation {
+        let world = self.world;
+
+        let location = world
+            .entities
+            .get(self.entity)
+            .expect("entity should exist at this point.");
+
+        let location = {
+            let archetype = &mut world.archetypes[location.archetype_id];
+            let ((disabled_arch, archetype_row), swapped_archetype) =
+                archetype.swap_disable(location.archetype_row);
+
+            // set the correct entity location for the swapped entity; the disabled is set to `None`
+            if let Some((entity, archetype_row)) = swapped_archetype {
+                let entity = entity.id();
+                let swapped_location = world.entities.get(entity).unwrap();
+
+                // SAFETY: TODO
+                unsafe {
+                    world.entities.set(
+                        entity.index(),
+                        Some(EntityLocation {
+                            archetype_row,
+                            ..swapped_location
+                        }),
+                    );
+                }
+            }
+
+            // SAFETY: TODO
+            let ((disabled_table, table_row), swapped_table) = unsafe {
+                world.storages.tables[archetype.table_id()]
+                    .swap_disable_unchecked(disabled_arch.table_row())
+            };
+
+            if let Some((entity, table_row)) = swapped_table {
+                let swapped_location = world.entities.get(entity).unwrap();
+
+                // SAFETY: TODO
+                unsafe {
+                    world.entities.set(
+                        entity.index(),
+                        Some(EntityLocation {
+                            table_row,
+                            ..swapped_location
+                        }),
+                    );
+                }
+            }
+
+            assert_eq!(disabled_table, disabled_arch.id());
+
+            EntityLocation {
+                archetype_row,
+                table_row,
+                ..location
+            }
+        };
+
+        // SAFETY: TODO
+        unsafe {
+            world.entities.set(self.entity.index(), None);
+        }
+
+        location
+    }
+
     pub(crate) fn despawn_with_caller(self, caller: MaybeLocation) {
         let location = self.location();
         let world = self.world;

crates/bevy_ecs/src/world/mod.rs

@@ -41,7 +41,7 @@ use crate::{
         ComponentTicks, Components, ComponentsQueuedRegistrator, ComponentsRegistrator, Mutable,
         RequiredComponents, RequiredComponentsError, Tick,
     },
-    entity::{Entities, Entity, EntityDoesNotExistError},
+    entity::{Entities, Entity, EntityDoesNotExistError, EntityLocation},
     entity_disabling::DefaultQueryFilters,
     error::{DefaultErrorHandler, ErrorHandler},
     lifecycle::{ComponentHooks, RemovedComponentMessages, ADD, DESPAWN, INSERT, REMOVE, REPLACE},
@@ -1439,6 +1439,13 @@ impl World {
         Ok(())
     }
 
+    pub fn disable(&mut self, entity: Entity) -> Result<EntityLocation, EntityMutableFetchError> {
+        self.flush();
+
+        let entity = self.get_entity_mut(entity)?;
+        Ok(entity.disable())
+    }
+
     /// Clears the internal component tracker state.
     ///
     /// The world maintains some internal state about changed and removed components. This state