[Merged by Bors] - Make EntityRef::new unsafe

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.

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<EntityMut<'w>> 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,

crates/bevy_ecs/src/world/mod.rs

@@ -317,7 +317,10 @@ impl World {
     #[inline]
     pub fn get_entity(&self, entity: Entity) -> Option<EntityRef> {
         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) }
                 })
         })
     }