Construct and deconstruct entities to improve entity allocation

crates/bevy_diagnostic/src/entity_count_diagnostics_plugin.rs

@@ -24,6 +24,6 @@ impl EntityCountDiagnosticsPlugin {
 
     /// Updates entity count measurement.
     pub fn diagnostic_system(mut diagnostics: Diagnostics, entities: &Entities) {
-        diagnostics.add_measurement(&Self::ENTITY_COUNT, || entities.len() as f64);
+        diagnostics.add_measurement(&Self::ENTITY_COUNT, || entities.count_constructed() as f64);
     }
 }

crates/bevy_ecs/src/bundle.rs

@@ -65,7 +65,7 @@ use crate::{
         Component, ComponentId, Components, ComponentsRegistrator, RequiredComponentConstructor,
         RequiredComponents, StorageType, Tick,
     },
-    entity::{Entities, Entity, EntityLocation},
+    entity::{Entities, EntitiesAllocator, Entity, EntityLocation},
     lifecycle::{ADD, INSERT, REMOVE, REPLACE},
     observer::Observers,
     prelude::World,
@@ -1221,9 +1221,9 @@ impl<'w> BundleInserter<'w> {
                 if let Some(swapped_entity) = result.swapped_entity {
                     let swapped_location =
                         // SAFETY: If the swap was successful, swapped_entity must be valid.
-                        unsafe { entities.get(swapped_entity).debug_checked_unwrap() };
-                    entities.set(
-                        swapped_entity.index(),
+                        unsafe { entities.get_constructed(swapped_entity).debug_checked_unwrap() };
+                    entities.update(
+                        swapped_entity.row(),
                         Some(EntityLocation {
                             archetype_id: swapped_location.archetype_id,
                             archetype_row: location.archetype_row,
@@ -1233,7 +1233,7 @@ impl<'w> BundleInserter<'w> {
                     );
                 }
                 let new_location = new_archetype.allocate(entity, result.table_row);
-                entities.set(entity.index(), Some(new_location));
+                entities.update(entity.row(), Some(new_location));
                 let after_effect = bundle_info.write_components(
                     table,
                     sparse_sets,
@@ -1270,9 +1270,9 @@ impl<'w> BundleInserter<'w> {
                 if let Some(swapped_entity) = result.swapped_entity {
                     let swapped_location =
                         // SAFETY: If the swap was successful, swapped_entity must be valid.
-                        unsafe { entities.get(swapped_entity).debug_checked_unwrap() };
-                    entities.set(
-                        swapped_entity.index(),
+                        unsafe { entities.get_constructed(swapped_entity).debug_checked_unwrap() };
+                    entities.update(
+                        swapped_entity.row(),
                         Some(EntityLocation {
                             archetype_id: swapped_location.archetype_id,
                             archetype_row: location.archetype_row,
@@ -1285,16 +1285,16 @@ impl<'w> BundleInserter<'w> {
                 // redundant copies
                 let move_result = table.move_to_superset_unchecked(result.table_row, new_table);
                 let new_location = new_archetype.allocate(entity, move_result.new_row);
-                entities.set(entity.index(), Some(new_location));
+                entities.update(entity.row(), Some(new_location));
 
                 // If an entity was moved into this entity's table spot, update its table row.
                 if let Some(swapped_entity) = move_result.swapped_entity {
                     let swapped_location =
                         // SAFETY: If the swap was successful, swapped_entity must be valid.
-                        unsafe { entities.get(swapped_entity).debug_checked_unwrap() };
+                        unsafe { entities.get_constructed(swapped_entity).debug_checked_unwrap() };
 
-                    entities.set(
-                        swapped_entity.index(),
+                    entities.update(
+                        swapped_entity.row(),
                         Some(EntityLocation {
                             archetype_id: swapped_location.archetype_id,
                             archetype_row: swapped_location.archetype_row,
@@ -1610,10 +1610,10 @@ impl<'w> BundleRemover<'w> {
             .swap_remove(location.archetype_row);
         // if an entity was moved into this entity's archetype row, update its archetype row
         if let Some(swapped_entity) = remove_result.swapped_entity {
-            let swapped_location = world.entities.get(swapped_entity).unwrap();
+            let swapped_location = world.entities.get_constructed(swapped_entity).unwrap();
 
-            world.entities.set(
-                swapped_entity.index(),
+            world.entities.update(
+                swapped_entity.row(),
                 Some(EntityLocation {
                     archetype_id: swapped_location.archetype_id,
                     archetype_row: location.archetype_row,
@@ -1651,10 +1651,10 @@ impl<'w> BundleRemover<'w> {
 
             // if an entity was moved into this entity's table row, update its table row
             if let Some(swapped_entity) = move_result.swapped_entity {
-                let swapped_location = world.entities.get(swapped_entity).unwrap();
+                let swapped_location = world.entities.get_constructed(swapped_entity).unwrap();
 
-                world.entities.set(
-                    swapped_entity.index(),
+                world.entities.update(
+                    swapped_entity.row(),
                     Some(EntityLocation {
                         archetype_id: swapped_location.archetype_id,
                         archetype_row: swapped_location.archetype_row,
@@ -1676,7 +1676,7 @@ impl<'w> BundleRemover<'w> {
 
         // SAFETY: The entity is valid and has been moved to the new location already.
         unsafe {
-            world.entities.set(entity.index(), Some(new_location));
+            world.entities.update(entity.row(), Some(new_location));
         }
 
         (new_location, pre_remove_result)
@@ -1751,10 +1751,14 @@ impl<'w> BundleSpawner<'w> {
     }
 
     /// # Safety
-    /// `entity` must be allocated (but non-existent), `T` must match this [`BundleInfo`]'s type
+    ///  `T` must match this [`BundleInfo`]'s type
+    ///
+    /// # Panics
+    ///
+    /// Panics if the entity has already been constructed.
     #[inline]
     #[track_caller]
-    pub unsafe fn spawn_non_existent<T: DynamicBundle>(
+    pub unsafe fn construct<T: DynamicBundle>(
         &mut self,
         entity: Entity,
         bundle: T,
@@ -1785,8 +1789,17 @@ impl<'w> BundleSpawner<'w> {
                 InsertMode::Replace,
                 caller,
             );
-            entities.set(entity.index(), Some(location));
-            entities.mark_spawn_despawn(entity.index(), caller, self.change_tick);
+
+            let was_at = entities.declare(entity.row(), Some(location));
+            // We need to assert here.
+            // Even if we can ensure that this entity is fresh from an allocator,
+            // it does not prevent users constructing arbitrary rows, which may overlap with the allocator.
+            // One alternative would be making `Entity` creation unsafe, but this is a good safety net anyway.
+            assert!(
+                was_at.is_none(),
+                "Constructing an {entity} that was already constructed is not allowed."
+            );
+            entities.mark_construct_or_destruct(entity.row(), caller, self.change_tick);
             (location, after_effect)
         };
 
@@ -1831,6 +1844,12 @@ impl<'w> BundleSpawner<'w> {
         (location, after_effect)
     }
 
+    #[inline]
+    pub fn allocator(&self) -> &'w EntitiesAllocator {
+        // SAFETY: No outstanding references to self.world, changes to allocator cannot invalidate our internal pointers
+        unsafe { &self.world.world().allocator }
+    }
+
     /// # Safety
     /// `T` must match this [`BundleInfo`]'s type
     #[inline]
@@ -1839,18 +1858,12 @@ impl<'w> BundleSpawner<'w> {
         bundle: T,
         caller: MaybeLocation,
     ) -> (Entity, T::Effect) {
-        let entity = self.entities().alloc();
+        let entity = self.allocator().alloc();
         // SAFETY: entity is allocated (but non-existent), `T` matches this BundleInfo's type
-        let (_, after_effect) = unsafe { self.spawn_non_existent(entity, bundle, caller) };
+        let (_, after_effect) = unsafe { self.construct(entity, bundle, caller) };
         (entity, after_effect)
     }
 
-    #[inline]
-    pub(crate) fn entities(&mut self) -> &mut Entities {
-        // SAFETY: No outstanding references to self.world, changes to entities cannot invalidate our internal pointers
-        unsafe { &mut self.world.world_mut().entities }
-    }
-
     /// # Safety
     /// - `Self` must be dropped after running this function as it may invalidate internal pointers.
     #[inline]

crates/bevy_ecs/src/entity/clone_entities.rs

@@ -1,20 +1,21 @@
-use alloc::{boxed::Box, collections::VecDeque, vec::Vec};
-use bevy_platform::collections::{hash_map::Entry, HashMap, HashSet};
-use bevy_ptr::{Ptr, PtrMut};
-use bevy_utils::prelude::DebugName;
-use bumpalo::Bump;
-use core::{any::TypeId, cell::LazyCell, ops::Range};
-use derive_more::derive::From;
-
 use crate::{
     archetype::Archetype,
     bundle::{Bundle, BundleId, InsertMode},
     component::{Component, ComponentCloneBehavior, ComponentCloneFn, ComponentId, ComponentInfo},
-    entity::{hash_map::EntityHashMap, Entities, Entity, EntityMapper},
+    entity::{hash_map::EntityHashMap, Entity, EntityMapper},
     query::DebugCheckedUnwrap,
     relationship::RelationshipHookMode,
     world::World,
 };
+use alloc::{boxed::Box, collections::VecDeque, vec::Vec};
+use bevy_platform::collections::{hash_map::Entry, HashMap, HashSet};
+use bevy_ptr::{Ptr, PtrMut};
+use bevy_utils::prelude::DebugName;
+use bumpalo::Bump;
+use core::{any::TypeId, cell::LazyCell, ops::Range};
+use derive_more::From;
+
+use super::EntitiesAllocator;
 
 /// Provides read access to the source component (the component being cloned) in a [`ComponentCloneFn`].
 pub struct SourceComponent<'a> {
@@ -78,7 +79,7 @@ pub struct ComponentCloneCtx<'a, 'b> {
     target_component_written: bool,
     bundle_scratch: &'a mut BundleScratch<'b>,
     bundle_scratch_allocator: &'b Bump,
-    entities: &'a Entities,
+    allocator: &'a EntitiesAllocator,
     source: Entity,
     target: Entity,
     component_info: &'a ComponentInfo,
@@ -104,7 +105,7 @@ impl<'a, 'b> ComponentCloneCtx<'a, 'b> {
         target: Entity,
         bundle_scratch_allocator: &'b Bump,
         bundle_scratch: &'a mut BundleScratch<'b>,
-        entities: &'a Entities,
+        allocator: &'a EntitiesAllocator,
         component_info: &'a ComponentInfo,
         entity_cloner: &'a mut EntityClonerState,
         mapper: &'a mut dyn EntityMapper,
@@ -118,7 +119,7 @@ impl<'a, 'b> ComponentCloneCtx<'a, 'b> {
             bundle_scratch,
             target_component_written: false,
             bundle_scratch_allocator,
-            entities,
+            allocator,
             mapper,
             component_info,
             state: entity_cloner,
@@ -271,7 +272,7 @@ impl<'a, 'b> ComponentCloneCtx<'a, 'b> {
 
     /// Queues the `entity` to be cloned by the current [`EntityCloner`]
     pub fn queue_entity_clone(&mut self, entity: Entity) {
-        let target = self.entities.reserve_entity();
+        let target = self.allocator.alloc();
         self.mapper.set_mapped(entity, target);
         self.state.clone_queue.push_back(entity);
     }
@@ -536,6 +537,10 @@ impl EntityCloner {
         relationship_hook_insert_mode: RelationshipHookMode,
     ) -> Entity {
         let target = mapper.get_mapped(source);
+        // The target may need to be constructed if it hasn't been already.
+        // If this fails, it either didn't need to be constructed (ok) or doesn't exist (caught better later).
+        let _ = world.construct_empty(target);
+
         // PERF: reusing allocated space across clones would be more efficient. Consider an allocation model similar to `Commands`.
         let bundle_scratch_allocator = Bump::new();
         let mut bundle_scratch: BundleScratch;
@@ -554,14 +559,17 @@ impl EntityCloner {
             #[cfg(not(feature = "bevy_reflect"))]
             let app_registry = Option::<()>::None;
 
-            let source_archetype = source_entity.archetype();
+            let source_archetype = source_entity
+                .archetype()
+                .expect("Source entity must exist constructed");
             bundle_scratch = BundleScratch::with_capacity(source_archetype.component_count());
 
             let target_archetype = LazyCell::new(|| {
                 world
                     .get_entity(target)
                     .expect("Target entity must exist")
                     .archetype()
+                    .expect("Target entity must exist constructed")
             });
 
             filter.clone_components(source_archetype, target_archetype, |component| {
@@ -598,7 +606,7 @@ impl EntityCloner {
                         target,
                         &bundle_scratch_allocator,
                         &mut bundle_scratch,
-                        world.entities(),
+                        world.entities_allocator(),
                         info,
                         state,
                         mapper,

crates/bevy_ecs/src/entity/map_entities.rs

@@ -56,7 +56,7 @@ use super::EntityIndexSet;
 pub trait MapEntities {
     /// Updates all [`Entity`] references stored inside using `entity_mapper`.
     ///
-    /// Implementors should look up any and all [`Entity`] values stored within `self` and
+    /// Implementers should look up any and all [`Entity`] values stored within `self` and
     /// update them to the mapped values via `entity_mapper`.
     fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E);
 }
@@ -202,7 +202,7 @@ impl<T: MapEntities, A: smallvec::Array<Item = T>> MapEntities for SmallVec<A> {
 ///
 /// More generally, this can be used to map [`Entity`] references between any two [`Worlds`](World).
 ///
-/// This is used by [`MapEntities`] implementors.
+/// This is used by [`MapEntities`] implementers.
 ///
 /// ## Example
 ///
@@ -336,14 +336,10 @@ impl<'m> SceneEntityMapper<'m> {
     }
 
     /// Creates a new [`SceneEntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
-    pub fn new(map: &'m mut EntityHashMap<Entity>, world: &mut World) -> Self {
-        // We're going to be calling methods on `Entities` that require advance
-        // flushing, such as `alloc` and `free`.
-        world.flush_entities();
+    pub fn new(map: &'m mut EntityHashMap<Entity>, world: &World) -> Self {
         Self {
             map,
-            // SAFETY: Entities data is kept in a valid state via `EntityMapper::world_scope`
-            dead_start: unsafe { world.entities_mut().alloc() },
+            dead_start: world.allocator.alloc(),
             generations: 0,
         }
     }
@@ -352,10 +348,13 @@ impl<'m> SceneEntityMapper<'m> {
     /// [`Entity`] while reserving extra generations. Because this makes the [`SceneEntityMapper`] unable to
     /// safely allocate any more references, this method takes ownership of `self` in order to render it unusable.
     pub fn finish(self, world: &mut World) {
-        // SAFETY: Entities data is kept in a valid state via `EntityMap::world_scope`
-        let entities = unsafe { world.entities_mut() };
-        assert!(entities.free(self.dead_start).is_some());
-        assert!(entities.reserve_generations(self.dead_start.index(), self.generations));
+        // SAFETY: We never constructed the entity and never released it for something else to construct.
+        let reuse_row = unsafe {
+            world
+                .entities
+                .mark_free(self.dead_start.row(), self.generations)
+        };
+        world.allocator.free(reuse_row);
     }
 
     /// Creates an [`SceneEntityMapper`] from a provided [`World`] and [`EntityHashMap<Entity>`], then calls the
@@ -388,7 +387,7 @@ mod tests {
     fn entity_mapper() {
         let mut map = EntityHashMap::default();
         let mut world = World::new();
-        let mut mapper = SceneEntityMapper::new(&mut map, &mut world);
+        let mut mapper = SceneEntityMapper::new(&mut map, &world);
 
         let mapped_ent = Entity::from_raw_u32(1).unwrap();
         let dead_ref = mapper.get_mapped(mapped_ent);
@@ -431,21 +430,4 @@ mod tests {
             .cmp_approx(&dead_ref.generation())
             .is_gt());
     }
-
-    #[test]
-    fn entity_mapper_no_panic() {
-        let mut world = World::new();
-        // "Dirty" the `Entities`, requiring a flush afterward.
-        world.entities.reserve_entity();
-        assert!(world.entities.needs_flush());
-
-        // Create and exercise a SceneEntityMapper - should not panic because it flushes
-        // `Entities` first.
-        SceneEntityMapper::world_scope(&mut Default::default(), &mut world, |_, m| {
-            m.get_mapped(Entity::PLACEHOLDER);
-        });
-
-        // The SceneEntityMapper should leave `Entities` in a flushed state.
-        assert!(!world.entities.needs_flush());
-    }
 }