component_scope

Author: janis-bhm

crates/bevy_ecs/src/archetype.rs

@@ -395,7 +395,7 @@ pub struct Archetype {
     table_id: TableId,
     edges: Edges,
     entities: Vec<ArchetypeEntity>,
-    disabled_entities: u32,
+    pub(crate) disabled_entities: u32,
     components: ImmutableSparseSet<ComponentId, ArchetypeComponentInfo>,
     pub(crate) flags: ArchetypeFlags,
 }
@@ -488,6 +488,12 @@ impl Archetype {
         &self.entities
     }
 
+    /// Fetches the disabled entities contained in this archetype.
+    #[inline]
+    pub fn disabled_entities(&self) -> &[ArchetypeEntity] {
+        &self.entities[..self.disabled_entities as usize]
+    }
+
     /// Get the valid table rows (i.e. non-disabled entities).
     #[inline]
     pub fn archetype_rows(&self) -> Range<u32> {
@@ -499,7 +505,10 @@ impl Archetype {
 
     /// Fetches the entities contained in this archetype.
     #[inline]
-    pub fn entities_with_location(&self) -> impl Iterator<Item = (Entity, EntityLocation)> {
+    pub fn entities_with_location(
+        &self,
+    ) -> impl Iterator<Item = (Entity, EntityLocation)> + DoubleEndedIterator + ExactSizeIterator
+    {
         self.entities.iter().enumerate().map(
             |(archetype_row, &ArchetypeEntity { entity, table_row })| {
                 (
@@ -635,9 +644,11 @@ 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.
+    /// Disables or enables the entity at `row` by swapping it with the first
+    /// enabled or disabled entity. Returns the swapped entities with their
+    /// respective table and archetype rows, or `None` if no swap occurred. If a
+    /// swap occurred, the caller is responsible for updating the entity's
+    /// location.
     ///
     /// # Panics
     /// This function will panic if `row >= self.entities.len()`
@@ -650,12 +661,29 @@ impl Archetype {
         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
+        let disabled_row = if row.index_u32() >= self.disabled_entities {
+            // the entity is currently enabled, swap it with the first enabled entity:
+
+            // 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.disabled_entities += 1;
 
+            disabled_row
+        } else {
+            self.disabled_entities -= 1;
+            // the entity is currently disabled, swap it with the last disabled entity:
+
+            // SAFETY: `self.disabled_entities` is always less than `u32::MAX`,
+            // as guaranteed by `allocate`.
+            ArchetypeRow::new(unsafe { NonMaxU32::new_unchecked(self.disabled_entities) })
+        };
+
+        if row == disabled_row {
+            // the entity is already in the correct position, no swap needed
+
             // SAFETY: `row` is guaranteed to be in-bounds.
             (
                 (unsafe { *self.entities.get_unchecked(row.index()) }, row),

crates/bevy_ecs/src/entity/mod.rs

@@ -1318,6 +1318,13 @@ pub struct EntityLocation {
 /// It is also useful for reserving an id; commands will often allocate an `Entity` but not provide it a location until the command is applied.
 pub type EntityIdLocation = Option<EntityLocation>;
 
+/// An [`Entity`] that has been disabled. This entity will not be found by queries or accessible via other ECS operations, but it's components are still stored in the ECS.
+/// This is useful for temporarily disabling an entity without fully despawning it or invoking archetype moves. This entity keeps track of its location, so it can be re-enabled later.
+pub struct DisabledEntity {
+    pub entity: Entity,
+    pub location: EntityLocation,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/bevy_ecs/src/query/iter.rs

@@ -2746,19 +2746,17 @@ mod tests {
             let mut query = world.query::<&A>();
             let mut iter = query.iter(&world);
             println!(
-                "archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             _ = iter.next();
             println!(
-                "archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             println!("{}", iter.sort::<Entity>().len());
         }
@@ -2776,19 +2774,17 @@ mod tests {
             let mut query = world.query::<&Sparse>();
             let mut iter = query.iter(&world);
             println!(
-                "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "before_next_call: archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             _ = iter.next();
             println!(
-                "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "after_next_call: archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             println!("{}", iter.sort::<Entity>().len());
         }
@@ -2801,19 +2797,17 @@ mod tests {
             let mut query = world.query::<(&A, &Sparse)>();
             let mut iter = query.iter(&world);
             println!(
-                "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "before_next_call: archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             _ = iter.next();
             println!(
-                "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "before_next_call: archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             println!("{}", iter.sort::<Entity>().len());
         }
@@ -2829,20 +2823,18 @@ mod tests {
             let mut query = world.query::<(&A, &Sparse)>();
             let mut iter = query.iter(&world);
             println!(
-                "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "before_next_call: archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             assert!(iter.cursor.table_entities.len() | iter.cursor.archetype_entities.len() == 0);
             _ = iter.next();
             println!(
-                "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
+                "before_next_call: archetype_entities: {} table_entities: {} current_range: {:?}",
                 iter.cursor.archetype_entities.len(),
                 iter.cursor.table_entities.len(),
-                iter.cursor.current_len,
-                iter.cursor.current_row
+                iter.cursor.current_range,
             );
             assert!(
                 (

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

@@ -198,7 +198,7 @@ impl Table {
     /// Fetches a read-only slice of the entities stored within the [`Table`].
     #[inline]
     pub fn entities(&self) -> &[Entity] {
-        &self.entities[self.disabled_entities as usize..]
+        &self.entities
     }
 
     /// Get the valid table rows (i.e. non-disabled entities).
@@ -217,7 +217,11 @@ 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`.
+    /// Disables or enables the entity at `row` by swapping it with the first
+    /// enabled or disabled entity. Returns the swapped entities with their
+    /// respective table rows, or `None` if no swap occurred. If a
+    /// swap occurred, the caller is responsible for updating the entity's
+    /// location.
     ///
     /// # Safety
     /// `row` must be in-bounds (`row.as_usize()` < `self.len()`) and not disabled.
@@ -226,20 +230,37 @@ impl Table {
         row: 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() != self.disabled_entities {
+        let disabled_row = if row.index_u32() >= self.disabled_entities {
+            // the entity is currently enabled, swap it with the first enabled entity:
+
             // SAFETY: `self.disabled_entities` is always less than `u32::MAX`,
             // as guaranteed by `allocate`.
             let disabled_row =
                 TableRow::new(unsafe { NonMaxU32::new_unchecked(self.disabled_entities) });
+            self.disabled_entities += 1;
+
+            disabled_row
+        } else {
+            self.disabled_entities -= 1;
+            // the entity is currently disabled, swap it with the last disabled entity:
+
+            // SAFETY: `self.disabled_entities` is always less than `u32::MAX`,
+            // as guaranteed by `allocate`.
+            TableRow::new(unsafe { NonMaxU32::new_unchecked(self.disabled_entities) })
+        };
 
+        if row == disabled_row {
+            // the entity is already in the correct position, no swap needed
+
+            // SAFETY: TODO
+            ((*self.entities.get_unchecked(row.index()), row), None)
+        } else {
             for col in self.columns.values_mut() {
                 col.swap_unchecked(row, disabled_row);
             }
 
             self.entities.swap(row.index(), disabled_row.index());
-            self.disabled_entities += 1;
 
             (
                 (
@@ -248,9 +269,6 @@ impl Table {
                 ),
                 Some((*self.entities.get_unchecked(row.index()), row)),
             )
-        } else {
-            self.disabled_entities += 1;
-            ((*self.entities.get_unchecked(row.index()), row), None)
         }
     }
 

crates/bevy_ecs/src/world/entity_ref.rs

@@ -6,8 +6,8 @@ use crate::{
     change_detection::{MaybeLocation, MutUntyped},
     component::{Component, ComponentId, ComponentTicks, Components, Mutable, StorageType, Tick},
     entity::{
-        ContainsEntity, Entity, EntityCloner, EntityClonerBuilder, EntityEquivalent,
-        EntityIdLocation, EntityLocation, OptIn, OptOut,
+        ContainsEntity, DisabledEntity, Entity, EntityCloner, EntityClonerBuilder,
+        EntityEquivalent, EntityIdLocation, EntityLocation, OptIn, OptOut,
     },
     event::{EntityComponentsTrigger, EntityEvent},
     lifecycle::{Despawn, Remove, Replace, DESPAWN, REMOVE, REPLACE},
@@ -2564,7 +2564,7 @@ impl<'w> EntityWorldMut<'w> {
 
     /// Disable the current entity.
     ///
-    pub fn disable(self) -> EntityLocation {
+    pub fn disable(self) -> DisabledEntity {
         let world = self.world;
 
         let location = world
@@ -2629,7 +2629,10 @@ impl<'w> EntityWorldMut<'w> {
             world.entities.set(self.entity.index(), None);
         }
 
-        location
+        DisabledEntity {
+            entity: self.entity,
+            location,
+        }
     }
 
     pub(crate) fn despawn_with_caller(self, caller: MaybeLocation) {

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, EntityLocation},
+    entity::{DisabledEntity, Entities, Entity, EntityDoesNotExistError},
     entity_disabling::DefaultQueryFilters,
     error::{DefaultErrorHandler, ErrorHandler},
     lifecycle::{ComponentHooks, RemovedComponentMessages, ADD, DESPAWN, INSERT, REMOVE, REPLACE},
@@ -1439,13 +1439,91 @@ impl World {
         Ok(())
     }
 
-    pub fn disable(&mut self, entity: Entity) -> Result<EntityLocation, EntityMutableFetchError> {
+    /// Disables the given `entity`, returning a [`DisabledEntity`] if the
+    /// entity exists and was enabled.
+    pub fn disable(&mut self, entity: Entity) -> Result<DisabledEntity, EntityMutableFetchError> {
         self.flush();
 
         let entity = self.get_entity_mut(entity)?;
         Ok(entity.disable())
     }
 
+    /// Re-enables a previously disabled entity, returning an [`EntityWorldMut`]
+    /// to it.
+    pub fn enable(&mut self, disabled: DisabledEntity) -> EntityWorldMut<'_> {
+        self.flush();
+
+        let archetype = &self.archetypes[disabled.location.archetype_id];
+        // Find the location of the disabled entity in the archetype by
+        // searching backwards through the disabled entities.
+        // The disabled entity is most likely the last disabled entity, as
+        // is the case when disabled for `entity_scope`.
+        let location = archetype
+            .entities_with_location()
+            .take(archetype.disabled_entities as usize)
+            .rev()
+            .find(|(e, _)| *e == disabled.entity)
+            .map(|(_, location)| location);
+
+        // SAFETY: disabled entity existed when it was disabled, and wasn't
+        // despawned in the meantime.
+        unsafe {
+            self.entities.set(disabled.entity.index(), location);
+        }
+
+        self.entity_mut(disabled.entity)
+    }
+
+    pub fn component_scope<
+        C: Component<Mutability = Mutable>,
+        R,
+        F: FnOnce(Entity, &mut C, &mut World) -> R,
+    >(
+        &mut self,
+        entity: Entity,
+        f: F,
+    ) {
+        self.try_component_scope(entity, f).unwrap_or_else(|| {
+            panic!(
+                "component {} does not exist for entity {entity}",
+                DebugName::type_name::<C>()
+            )
+        });
+    }
+
+    pub fn try_component_scope<
+        C: Component<Mutability = Mutable>,
+        R,
+        F: FnOnce(Entity, &mut C, &mut World) -> R,
+    >(
+        &mut self,
+        entity: Entity,
+        f: F,
+    ) -> Option<R> {
+        self.flush();
+
+        let mut entity_mut = self.get_entity_mut(entity).ok()?;
+        let mut component = {
+            let component = entity_mut.get_mut::<C>()?;
+            // SAFETY: TODO
+            unsafe { core::ptr::read::<C>(&raw const *component) }
+        };
+
+        let disabled = entity_mut.disable();
+
+        let out = f(entity, &mut component, self);
+
+        let mut entity_mut = self.enable(disabled);
+
+        let mut component_mut = entity_mut.get_mut::<C>().unwrap();
+        // SAFETY: TODO
+        unsafe {
+            core::ptr::write::<C>(&raw mut *component_mut, component);
+        }
+
+        Some(out)
+    }
+
     /// Clears the internal component tracker state.
     ///
     /// The world maintains some internal state about changed and removed components. This state