diff --git a/crates/bevy_ecs/src/entity/mod.rs b/crates/bevy_ecs/src/entity/mod.rs index 5107923e4a..afcba27bb6 100644 --- a/crates/bevy_ecs/src/entity/mod.rs +++ b/crates/bevy_ecs/src/entity/mod.rs @@ -747,7 +747,7 @@ impl EntityMeta { // SAFETY: // This type must not contain any pointers at any level, and be safe to fully fill with u8::MAX. /// A location of an entity in an archetype. -#[derive(Copy, Clone, Debug)] +#[derive(Copy, Clone, Debug, PartialEq)] #[repr(C)] pub struct EntityLocation { /// The ID of the [`Archetype`] the [`Entity`] belongs to. diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index f21675103f..a165678516 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -22,8 +22,17 @@ pub struct EntityRef<'w> { } impl<'w> EntityRef<'w> { + /// # Safety + /// + /// - `entity` must be valid for `world`: the generation should match that of the entity at the same index. + /// - `location` must be sourced from `world`'s `Entities` and must exactly match the location for `entity` + /// + /// The above is trivially satisfied if `location` was sourced from `world.entities().get(entity)`. #[inline] - pub(crate) fn new(world: &'w World, entity: Entity, location: EntityLocation) -> Self { + pub(crate) unsafe fn new(world: &'w World, entity: Entity, location: EntityLocation) -> Self { + debug_assert!(world.entities().contains(entity)); + debug_assert_eq!(world.entities().get(entity), Some(location)); + Self { world, entity, @@ -193,7 +202,9 @@ impl<'w> EntityRef<'w> { impl<'w> From> for EntityRef<'w> { fn from(entity_mut: EntityMut<'w>) -> EntityRef<'w> { - EntityRef::new(entity_mut.world, entity_mut.entity, entity_mut.location) + // SAFETY: the safety invariants on EntityMut and EntityRef are identical + // and EntityMut is promised to be valid by construction. + unsafe { EntityRef::new(entity_mut.world, entity_mut.entity, entity_mut.location) } } } @@ -206,13 +217,20 @@ pub struct EntityMut<'w> { impl<'w> EntityMut<'w> { /// # Safety - /// entity and location _must_ be valid + /// + /// - `entity` must be valid for `world`: the generation should match that of the entity at the same index. + /// - `location` must be sourced from `world`'s `Entities` and must exactly match the location for `entity` + /// + /// The above is trivially satisfied if `location` was sourced from `world.entities().get(entity)`. #[inline] pub(crate) unsafe fn new( world: &'w mut World, entity: Entity, location: EntityLocation, ) -> Self { + debug_assert!(world.entities().contains(entity)); + debug_assert_eq!(world.entities().get(entity), Some(location)); + EntityMut { world, entity, diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 2243af65de..af48073446 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -317,7 +317,10 @@ impl World { #[inline] pub fn get_entity(&self, entity: Entity) -> Option { let location = self.entities.get(entity)?; - Some(EntityRef::new(self, entity, location)) + // SAFETY: if the Entity is invalid, the function returns early. + // Additionally, Entities::get(entity) returns the correct EntityLocation if the entity exists. + let entity_ref = unsafe { EntityRef::new(self, entity, location) }; + Some(entity_ref) } /// Returns an [`Entity`] iterator of current entities. @@ -331,13 +334,16 @@ impl World { .iter() .enumerate() .map(|(archetype_row, archetype_entity)| { + let entity = archetype_entity.entity(); let location = EntityLocation { archetype_id: archetype.id(), archetype_row: ArchetypeRow::new(archetype_row), table_id: archetype.table_id(), table_row: archetype_entity.table_row(), }; - EntityRef::new(self, archetype_entity.entity(), location) + + // SAFETY: entity exists and location accurately specifies the archetype where the entity is stored + unsafe { EntityRef::new(self, entity, location) } }) }) }