Skip to content

Commit

Permalink
Lock down access to Entities (#6740)
Browse files Browse the repository at this point in the history
# 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 #3362 for `bevy_ecs::entity`. Incorporates the changes from #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 <contact@jamessliu.com>
Co-authored-by: James Liu <contact@jamessliu.com>
  • Loading branch information
nakedible and james7132 committed Nov 28, 2022
1 parent 9c79b39 commit e954b85
Show file tree
Hide file tree
Showing 7 changed files with 126 additions and 65 deletions.
10 changes: 5 additions & 5 deletions crates/bevy_ecs/src/bundle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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
}
Expand Down
36 changes: 36 additions & 0 deletions crates/bevy_ecs/src/entity/map_entities.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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>;
}

Expand Down
76 changes: 63 additions & 13 deletions crates/bevy_ecs/src/entity/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
///
Expand Down Expand Up @@ -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:
Expand All @@ -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:
///
Expand Down Expand Up @@ -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`.
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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();
Expand All @@ -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
///
Expand Down Expand Up @@ -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
Expand All @@ -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,
}
Expand Down Expand Up @@ -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(|_, _| {}) };
Expand All @@ -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());
Expand Down
47 changes: 10 additions & 37 deletions crates/bevy_ecs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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());
Expand Down Expand Up @@ -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"
Expand All @@ -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()
Expand All @@ -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"
Expand All @@ -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"
);
Expand Down Expand Up @@ -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))];

Expand Down
11 changes: 8 additions & 3 deletions crates/bevy_ecs/src/world/entity_ref.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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;

Expand Down
Loading

0 comments on commit e954b85

Please sign in to comment.