Lock down access to Entities (bevyengine#6740)

# Objective
The soundness of the ECS `World` partially relies on the correctness of the state of `Entities` stored within it. We're currently allowing users to (unsafely) mutate it, as well as readily construct it without using a `World`. While this is not strictly unsound so long as users (including `bevy_render`) safely use the APIs, it's a fairly easy path to unsoundness without much of a guard rail.

Addresses bevyengine#3362 for `bevy_ecs::entity`. Incorporates the changes from bevyengine#3985.

## Solution
Remove `Entities`'s  `Default` implementation and force access to the type to only be through a properly constructed `World`.

Additional cleanup for other parts of `bevy_ecs::entity`:

 - `Entity::index` and `Entity::generation` are no longer `pub(crate)`, opting to force the rest of bevy_ecs to use the public interface to access these values.
 - `EntityMeta` is no longer `pub` and also not `pub(crate)` to attempt to cut down on updating `generation` without going through an `Entities` API. It's currently inaccessible except via the `pub(crate)` Vec on `Entities`, there was no way for an outside user to use it.
 - Added `Entities::set`, an unsafe `pub(crate)` API for setting the location of an Entity (parallel to `Entities::get`) that replaces the internal case where we need to set the location of an entity when it's been spawned, moved, or despawned.
 - `Entities::alloc_at_without_replacement` is only used in `World::get_or_spawn` within the first party crates, and I cannot find a public use of this API in any ecosystem crate that I've checked (via GitHub search).
 - Attempted to document the few remaining undocumented public APIs in the module.

---

## Changelog
Removed: `Entities`'s `Default` implementation.
Removed: `EntityMeta`
Removed: `Entities::alloc_at_without_replacement` and `AllocAtWithoutReplacement`.

Co-authored-by: james7132 <[email protected]>
Co-authored-by: James Liu <[email protected]>

crates/bevy_ecs/src/bundle.rs

@@ -553,10 +553,10 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
             InsertBundleResult::NewArchetypeSameTable { new_archetype } => {
                 let result = self.archetype.swap_remove(location.index);
                 if let Some(swapped_entity) = result.swapped_entity {
-                    self.entities.meta[swapped_entity.index as usize].location = location;
+                    self.entities.set(swapped_entity.index(), location);
                 }
                 let new_location = new_archetype.allocate(entity, result.table_row);
-                self.entities.meta[entity.index as usize].location = new_location;
+                self.entities.set(entity.index(), new_location);
 
                 // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty)
                 let add_bundle = self
@@ -581,15 +581,15 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
             } => {
                 let result = self.archetype.swap_remove(location.index);
                 if let Some(swapped_entity) = result.swapped_entity {
-                    self.entities.meta[swapped_entity.index as usize].location = location;
+                    self.entities.set(swapped_entity.index(), location);
                 }
                 // PERF: store "non bundle" components in edge, then just move those to avoid
                 // redundant copies
                 let move_result = self
                     .table
                     .move_to_superset_unchecked(result.table_row, new_table);
                 let new_location = new_archetype.allocate(entity, move_result.new_row);
-                self.entities.meta[entity.index as usize].location = new_location;
+                self.entities.set(entity.index(), 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 {
@@ -664,7 +664,7 @@ impl<'a, 'b> BundleSpawner<'a, 'b> {
             self.change_tick,
             bundle,
         );
-        self.entities.meta[entity.index as usize].location = location;
+        self.entities.set(entity.index(), location);
 
         location
     }

crates/bevy_ecs/src/entity/map_entities.rs

@@ -2,6 +2,7 @@ use crate::entity::Entity;
 use bevy_utils::{Entry, HashMap};
 use std::fmt;
 
+/// The errors that might be returned while using [`MapEntities::map_entities`].
 #[derive(Debug)]
 pub enum MapEntitiesError {
     EntityNotFound(Entity),
@@ -19,7 +20,42 @@ impl fmt::Display for MapEntitiesError {
     }
 }
 
+/// Operation to map all contained [`Entity`] fields in a type to new values.
+///
+/// As entity IDs are valid only for the [`World`] they're sourced from, using [`Entity`]
+/// as references in components copied from another world will be invalid. This trait
+/// allows defining custom mappings for these references via [`EntityMap`].
+///
+/// Implementing this trait correctly is required for properly loading components
+/// with entity references from scenes.
+///
+/// ## Example
+///
+/// ```rust
+/// use bevy_ecs::prelude::*;
+/// use bevy_ecs::entity::{EntityMap, MapEntities, MapEntitiesError};
+///
+/// #[derive(Component)]
+/// struct Spring {
+///     a: Entity,
+///     b: Entity,
+/// }
+///
+/// impl MapEntities for Spring {
+///     fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError> {
+///         self.a = entity_map.get(self.a)?;
+///         self.b = entity_map.get(self.b)?;
+///         Ok(())
+///     }
+/// }
+/// ```
+///
+/// [`World`]: crate::world::World
 pub trait MapEntities {
+    /// Updates all [`Entity`] references stored inside using `entity_map`.
+    ///
+    /// Implementors should look up any and all [`Entity`] values stored within and
+    /// update them to the mapped values via `entity_map`.
     fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError>;
 }
 

crates/bevy_ecs/src/entity/mod.rs

@@ -56,6 +56,9 @@ type IdCursor = isize;
 /// The identifier is implemented using a [generational index]: a combination of an index and a generation.
 /// This allows fast insertion after data removal in an array while minimizing loss of spatial locality.
 ///
+/// These identifiers are only valid on the [`World`] it's sourced from. Attempting to use an `Entity` to
+/// fetch entity components or metadata from a different world will either fail or return unexpected results.
+///
 /// [generational index]: https://lucassardois.medium.com/generational-indices-guide-8e3c5f7fd594
 ///
 /// # Usage
@@ -103,19 +106,25 @@ type IdCursor = isize;
 /// [`EntityMut::id`]: crate::world::EntityMut::id
 /// [`EntityCommands`]: crate::system::EntityCommands
 /// [`Query::get`]: crate::system::Query::get
+/// [`World`]: crate::world::World
 #[derive(Clone, Copy, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct Entity {
-    pub(crate) generation: u32,
-    pub(crate) index: u32,
+    generation: u32,
+    index: u32,
 }
 
-pub enum AllocAtWithoutReplacement {
+pub(crate) enum AllocAtWithoutReplacement {
     Exists(EntityLocation),
     DidNotExist,
     ExistsWithWrongGeneration,
 }
 
 impl Entity {
+    #[cfg(test)]
+    pub(crate) const fn new(index: u32, generation: u32) -> Entity {
+        Entity { index, generation }
+    }
+
     /// An entity ID with a placeholder value. This may or may not correspond to an actual entity,
     /// and should be overwritten by a new value before being used.
     ///
@@ -266,9 +275,17 @@ impl<'a> Iterator for ReserveEntitiesIterator<'a> {
 impl<'a> core::iter::ExactSizeIterator for ReserveEntitiesIterator<'a> {}
 impl<'a> core::iter::FusedIterator for ReserveEntitiesIterator<'a> {}
 
-#[derive(Debug, Default)]
+/// A [`World`]'s internal metadata store on all of its entities.
+///
+/// Contains metadata on:
+///  - The generation of every entity.
+///  - The alive/dead status of a particular entity. (i.e. "has entity 3 been despawned?")
+///  - The location of the entity's components in memory (via [`EntityLocation`])
+///
+/// [`World`]: crate::world::World
+#[derive(Debug)]
 pub struct Entities {
-    pub(crate) meta: Vec<EntityMeta>,
+    meta: Vec<EntityMeta>,
 
     /// The `pending` and `free_cursor` fields describe three sets of Entity IDs
     /// that have been freed or are in the process of being allocated:
@@ -281,8 +298,8 @@ pub struct Entities {
     ///   `reserve_entities` or `reserve_entity()`. They are now waiting for `flush()` to make them
     ///   fully allocated.
     ///
-    /// - The count of new IDs that do not yet exist in `self.meta()`, but which we have handed out
-    ///   and reserved. `flush()` will allocate room for them in `self.meta()`.
+    /// - The count of new IDs that do not yet exist in `self.meta`, but which we have handed out
+    ///   and reserved. `flush()` will allocate room for them in `self.meta`.
     ///
     /// The contents of `pending` look like this:
     ///
@@ -312,6 +329,15 @@ pub struct Entities {
 }
 
 impl Entities {
+    pub(crate) const fn new() -> Self {
+        Entities {
+            meta: Vec::new(),
+            pending: Vec::new(),
+            free_cursor: AtomicIdCursor::new(0),
+            len: 0,
+        }
+    }
+
     /// Reserve entity IDs concurrently.
     ///
     /// Storage for entity generation and location is lazily allocated by calling `flush`.
@@ -448,7 +474,10 @@ impl Entities {
     /// Allocate a specific entity ID, overwriting its generation.
     ///
     /// Returns the location of the entity currently using the given ID, if any.
-    pub fn alloc_at_without_replacement(&mut self, entity: Entity) -> AllocAtWithoutReplacement {
+    pub(crate) fn alloc_at_without_replacement(
+        &mut self,
+        entity: Entity,
+    ) -> AllocAtWithoutReplacement {
         self.verify_flushed();
 
         let result = if entity.index as usize >= self.meta.len() {
@@ -523,6 +552,7 @@ impl Entities {
             .map_or(false, |e| e.generation() == entity.generation)
     }
 
+    /// Clears all [`Entity`] from the World.
     pub fn clear(&mut self) {
         self.meta.clear();
         self.pending.clear();
@@ -545,6 +575,18 @@ impl Entities {
         }
     }
 
+    /// Updates the location of an [`Entity`]. This must be called when moving the components of
+    /// the entity around in storage.
+    ///
+    /// # Safety
+    ///  - `index` must be a valid entity index.
+    ///  - `location` must be valid for the entity at `index` or immediately made valid afterwards
+    ///    before handing control to unknown code.
+    pub(crate) unsafe fn set(&mut self, index: u32, location: EntityLocation) {
+        // SAFETY: Caller guarentees that `index` a valid entity index
+        self.meta.get_unchecked_mut(index as usize).location = location;
+    }
+
     /// Get the [`Entity`] with a given id, if it exists in this [`Entities`] collection
     /// Returns `None` if this [`Entity`] is outside of the range of currently reserved Entities
     ///
@@ -646,17 +688,25 @@ impl Entities {
         self.len = count as u32;
     }
 
-    /// Accessor for getting the length of the vec in `self.meta`
+    /// The count of all entities in the [`World`] that have ever been allocated
+    /// including the entities that are currently freed.
+    ///
+    /// This does not include entities that have been reserved but have never been
+    /// allocated yet.
+    ///
+    /// [`World`]: crate::world::World
     #[inline]
-    pub fn meta_len(&self) -> usize {
+    pub fn total_count(&self) -> usize {
         self.meta.len()
     }
 
+    /// The count of currently allocated entities.
     #[inline]
     pub fn len(&self) -> u32 {
         self.len
     }
 
+    /// Checks if any entity is currently active.
     #[inline]
     pub fn is_empty(&self) -> bool {
         self.len == 0
@@ -667,7 +717,7 @@ impl Entities {
 // This type must not contain any pointers at any level, and be safe to fully fill with u8::MAX.
 #[derive(Copy, Clone, Debug)]
 #[repr(C)]
-pub struct EntityMeta {
+struct EntityMeta {
     pub generation: u32,
     pub location: EntityLocation,
 }
@@ -708,7 +758,7 @@ mod tests {
 
     #[test]
     fn reserve_entity_len() {
-        let mut e = Entities::default();
+        let mut e = Entities::new();
         e.reserve_entity();
         // SAFETY: entity_location is left invalid
         unsafe { e.flush(|_, _| {}) };
@@ -717,7 +767,7 @@ mod tests {
 
     #[test]
     fn get_reserved_and_invalid() {
-        let mut entities = Entities::default();
+        let mut entities = Entities::new();
         let e = entities.reserve_entity();
         assert!(entities.contains(e));
         assert!(entities.get(e).is_none());

crates/bevy_ecs/src/lib.rs

@@ -1465,7 +1465,7 @@ mod tests {
         let e3 = world_a.entities().reserve_entity();
         world_a.flush();
 
-        let world_a_max_entities = world_a.entities().meta.len();
+        let world_a_max_entities = world_a.entities().len();
         world_b
             .entities
             .reserve_entities(world_a_max_entities as u32);
@@ -1474,10 +1474,7 @@ mod tests {
         let e4 = world_b.spawn(A(4)).id();
         assert_eq!(
             e4,
-            Entity {
-                generation: 0,
-                index: 3,
-            },
+            Entity::new(3, 0),
             "new entity is created immediately after world_a's max entity"
         );
         assert!(world_b.get::<A>(e1).is_none());
@@ -1508,10 +1505,7 @@ mod tests {
             "spawning into existing `world_b` entities works"
         );
 
-        let e4_mismatched_generation = Entity {
-            generation: 1,
-            index: 3,
-        };
+        let e4_mismatched_generation = Entity::new(3, 1);
         assert!(
             world_b.get_or_spawn(e4_mismatched_generation).is_none(),
             "attempting to spawn on top of an entity with a mismatched entity generation fails"
@@ -1527,10 +1521,7 @@ mod tests {
             "failed mismatched spawn doesn't change existing entity"
         );
 
-        let high_non_existent_entity = Entity {
-            generation: 0,
-            index: 6,
-        };
+        let high_non_existent_entity = Entity::new(6, 0);
         world_b
             .get_or_spawn(high_non_existent_entity)
             .unwrap()
@@ -1541,10 +1532,7 @@ mod tests {
             "inserting into newly allocated high / non-continous entity id works"
         );
 
-        let high_non_existent_but_reserved_entity = Entity {
-            generation: 0,
-            index: 5,
-        };
+        let high_non_existent_but_reserved_entity = Entity::new(5, 0);
         assert!(
             world_b.get_entity(high_non_existent_but_reserved_entity).is_none(),
             "entities between high-newly allocated entity and continuous block of existing entities don't exist"
@@ -1560,22 +1548,10 @@ mod tests {
         assert_eq!(
             reserved_entities,
             vec![
-                Entity {
-                    generation: 0,
-                    index: 5
-                },
-                Entity {
-                    generation: 0,
-                    index: 4
-                },
-                Entity {
-                    generation: 0,
-                    index: 7,
-                },
-                Entity {
-                    generation: 0,
-                    index: 8,
-                },
+                Entity::new(5, 0),
+                Entity::new(4, 0),
+                Entity::new(7, 0),
+                Entity::new(8, 0),
             ],
             "space between original entities and high entities is used for new entity ids"
         );
@@ -1624,10 +1600,7 @@ mod tests {
         let e0 = world.spawn(A(0)).id();
         let e1 = Entity::from_raw(1);
         let e2 = world.spawn_empty().id();
-        let invalid_e2 = Entity {
-            generation: 1,
-            index: e2.index,
-        };
+        let invalid_e2 = Entity::new(e2.index(), 1);
 
         let values = vec![(e0, (B(0), C)), (e1, (B(1), C)), (invalid_e2, (B(2), C))];
 

crates/bevy_ecs/src/world/entity_ref.rs

@@ -360,7 +360,7 @@ impl<'w> EntityMut<'w> {
         let old_archetype = &mut archetypes[old_archetype_id];
         let remove_result = old_archetype.swap_remove(old_location.index);
         if let Some(swapped_entity) = remove_result.swapped_entity {
-            entities.meta[swapped_entity.index as usize].location = old_location;
+            entities.set(swapped_entity.index(), old_location);
         }
         let old_table_row = remove_result.table_row;
         let old_table_id = old_archetype.table_id();
@@ -394,7 +394,8 @@ impl<'w> EntityMut<'w> {
         };
 
         *self_location = new_location;
-        entities.meta[entity.index as usize].location = new_location;
+        // SAFETY: The entity is valid and has been moved to the new location already.
+        entities.set(entity.index(), new_location);
     }
 
     #[deprecated(
@@ -490,7 +491,11 @@ impl<'w> EntityMut<'w> {
             }
             let remove_result = archetype.swap_remove(location.index);
             if let Some(swapped_entity) = remove_result.swapped_entity {
-                world.entities.meta[swapped_entity.index as usize].location = location;
+                // SAFETY: swapped_entity is valid and the swapped entity's components are
+                // moved to the new location immediately after.
+                unsafe {
+                    world.entities.set(swapped_entity.index(), location);
+                }
             }
             table_row = remove_result.table_row;