Change Entity::generation from u32 to NonZeroU32 for niche optimization (#9907)

# Objective

- Implements change described in
#3022
- Goal is to allow Entity to benefit from niche optimization, especially
in the case of Option<Entity> to reduce memory overhead with structures
with empty slots

## Discussion
- First PR attempt: #3029
- Discord:
https://discord.com/channels/691052431525675048/1154573759752183808/1154573764240093224

## Solution

- Change `Entity::generation` from u32 to NonZeroU32 to allow for niche
optimization.
- The reason for changing generation rather than index is so that the
costs are only encountered on Entity free, instead of on Entity alloc
- There was some concern with generations being used, due to there being
some desire to introduce flags. This was more to do with the original
retirement approach, however, in reality even if generations were
reduced to 24-bits, we would still have 16 million generations available
before wrapping and current ideas indicate that we would be using closer
to 4-bits for flags.
- Additionally, another concern was the representation of relationships
where NonZeroU32 prevents us using the full address space, talking with
Joy it seems unlikely to be an issue. The majority of the time these
entity references will be low-index entries (ie. `ChildOf`, `Owes`),
these will be able to be fast lookups, and the remainder of the range
can use slower lookups to map to the address space.
- It has the additional benefit of being less visible to most users,
since generation is only ever really set through `from_bits` type
methods.
- `EntityMeta` was changed to match
- On free, generation now explicitly wraps:
- Originally, generation would panic in debug mode and wrap in release
mode due to using regular ops.
- The first attempt at this PR changed the behavior to "retire" slots
and remove them from use when generations overflowed. This change was
controversial, and likely needs a proper RFC/discussion.
- Wrapping matches current release behaviour, and should therefore be
less controversial.
- Wrapping also more easily migrates to the retirement approach, as
users likely to exhaust the exorbitant supply of generations will code
defensively against aliasing and that defensive code is less likely to
break than code assuming that generations don't wrap.
- We use some unsafe code here when wrapping generations, to avoid
branch on NonZeroU32 construction. It's guaranteed safe due to how we
perform wrapping and it results in significantly smaller ASM code.
    - https://godbolt.org/z/6b6hj8PrM 

## Migration

- Previous `bevy_scene` serializations have a high likelihood of being
broken, as they contain 0th generation entities.

## Current Issues
 
- `Entities::reserve_generations` and `EntityMapper` wrap now, even in
debug - although they technically did in release mode already so this
probably isn't a huge issue. It just depends if we need to change
anything here?

---------

Co-authored-by: Natalie Baker <[email protected]>

Author: notverymoe

benches/benches/bevy_utils/entity_hash.rs

@@ -13,12 +13,12 @@ fn make_entity(rng: &mut impl Rng, size: usize) -> Entity {
     // -log₂(1-x) gives an exponential distribution with median 1.0
     // That lets us get values that are mostly small, but some are quite large
     // * For ids, half are in [0, size), half are unboundedly larger.
-    // * For generations, half are in [0, 2), half are unboundedly larger.
+    // * For generations, half are in [1, 3), half are unboundedly larger.
 
     let x: f64 = rng.gen();
     let id = -(1.0 - x).log2() * (size as f64);
     let x: f64 = rng.gen();
-    let gen = -(1.0 - x).log2() * 2.0;
+    let gen = 1.0 + -(1.0 - x).log2() * 2.0;
 
     // this is not reliable, but we're internal so a hack is ok
     let bits = ((gen as u64) << 32) | (id as u64);

crates/bevy_ecs/src/entity/map_entities.rs

@@ -1,6 +1,8 @@
 use crate::{entity::Entity, world::World};
 use bevy_utils::EntityHashMap;
 
+use super::inc_generation_by;
+
 /// 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`]
@@ -69,7 +71,7 @@ impl<'m> EntityMapper<'m> {
 
         // this new entity reference is specifically designed to never represent any living entity
         let new = Entity {
-            generation: self.dead_start.generation + self.generations,
+            generation: inc_generation_by(self.dead_start.generation, self.generations),
             index: self.dead_start.index,
         };
         self.generations += 1;
@@ -146,7 +148,7 @@ mod tests {
         let mut world = World::new();
         let mut mapper = EntityMapper::new(&mut map, &mut world);
 
-        let mapped_ent = Entity::new(FIRST_IDX, 0);
+        let mapped_ent = Entity::from_raw(FIRST_IDX);
         let dead_ref = mapper.get_or_reserve(mapped_ent);
 
         assert_eq!(
@@ -155,7 +157,7 @@ mod tests {
             "should persist the allocated mapping from the previous line"
         );
         assert_eq!(
-            mapper.get_or_reserve(Entity::new(SECOND_IDX, 0)).index(),
+            mapper.get_or_reserve(Entity::from_raw(SECOND_IDX)).index(),
             dead_ref.index(),
             "should re-use the same index for further dead refs"
         );
@@ -173,7 +175,7 @@ mod tests {
         let mut world = World::new();
 
         let dead_ref = EntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
-            mapper.get_or_reserve(Entity::new(0, 0))
+            mapper.get_or_reserve(Entity::from_raw(0))
         });
 
         // Next allocated entity should be a further generation on the same index

crates/bevy_ecs/src/entity/mod.rs

@@ -37,14 +37,15 @@
 //! [`EntityWorldMut::remove`]: crate::world::EntityWorldMut::remove
 mod map_entities;
 
+use bevy_utils::tracing::warn;
 pub use map_entities::*;
 
 use crate::{
     archetype::{ArchetypeId, ArchetypeRow},
     storage::{SparseSetIndex, TableId, TableRow},
 };
 use serde::{Deserialize, Serialize};
-use std::{convert::TryFrom, fmt, hash::Hash, mem, sync::atomic::Ordering};
+use std::{convert::TryFrom, fmt, hash::Hash, mem, num::NonZeroU32, sync::atomic::Ordering};
 
 #[cfg(target_has_atomic = "64")]
 use std::sync::atomic::AtomicI64 as AtomicIdCursor;
@@ -124,7 +125,7 @@ pub struct Entity {
     // to make this struct equivalent to a u64.
     #[cfg(target_endian = "little")]
     index: u32,
-    generation: u32,
+    generation: NonZeroU32,
     #[cfg(target_endian = "big")]
     index: u32,
 }
@@ -188,8 +189,12 @@ pub(crate) enum AllocAtWithoutReplacement {
 
 impl Entity {
     #[cfg(test)]
-    pub(crate) const fn new(index: u32, generation: u32) -> Entity {
-        Entity { index, generation }
+    pub(crate) const fn new(index: u32, generation: u32) -> Result<Entity, &'static str> {
+        if let Some(generation) = NonZeroU32::new(generation) {
+            Ok(Entity { index, generation })
+        } else {
+            Err("Failed to construct Entity, check that generation > 0")
+        }
     }
 
     /// An entity ID with a placeholder value. This may or may not correspond to an actual entity,
@@ -228,7 +233,7 @@ impl Entity {
     /// ```
     pub const PLACEHOLDER: Self = Self::from_raw(u32::MAX);
 
-    /// Creates a new entity ID with the specified `index` and a generation of 0.
+    /// Creates a new entity ID with the specified `index` and a generation of 1.
     ///
     /// # Note
     ///
@@ -244,7 +249,7 @@ impl Entity {
     pub const fn from_raw(index: u32) -> Entity {
         Entity {
             index,
-            generation: 0,
+            generation: NonZeroU32::MIN,
         }
     }
 
@@ -256,17 +261,41 @@ impl Entity {
     /// No particular structure is guaranteed for the returned bits.
     #[inline(always)]
     pub const fn to_bits(self) -> u64 {
-        (self.generation as u64) << 32 | self.index as u64
+        (self.generation.get() as u64) << 32 | self.index as u64
     }
 
     /// Reconstruct an `Entity` previously destructured with [`Entity::to_bits`].
     ///
     /// Only useful when applied to results from `to_bits` in the same instance of an application.
-    #[inline(always)]
+    ///
+    /// # Panics
+    ///
+    /// This method will likely panic if given `u64` values that did not come from [`Entity::to_bits`]
+    #[inline]
     pub const fn from_bits(bits: u64) -> Self {
-        Self {
-            generation: (bits >> 32) as u32,
-            index: bits as u32,
+        // Construct an Identifier initially to extract the kind from.
+        let id = Self::try_from_bits(bits);
+
+        match id {
+            Ok(entity) => entity,
+            Err(_) => panic!("Attempted to initialise invalid bits as an entity"),
+        }
+    }
+
+    /// Reconstruct an `Entity` previously destructured with [`Entity::to_bits`].
+    ///
+    /// Only useful when applied to results from `to_bits` in the same instance of an application.
+    ///
+    /// This method is the fallible counterpart to [`Entity::from_bits`].
+    #[inline(always)]
+    pub const fn try_from_bits(bits: u64) -> Result<Self, &'static str> {
+        if let Some(generation) = NonZeroU32::new((bits >> 32) as u32) {
+            Ok(Self {
+                generation,
+                index: bits as u32,
+            })
+        } else {
+            Err("Invalid generation bits")
         }
     }
 
@@ -285,7 +314,7 @@ impl Entity {
     /// given index has been reused (index, generation) pairs uniquely identify a given Entity.
     #[inline]
     pub const fn generation(self) -> u32 {
-        self.generation
+        self.generation.get()
     }
 }
 
@@ -303,8 +332,9 @@ impl<'de> Deserialize<'de> for Entity {
     where
         D: serde::Deserializer<'de>,
     {
+        use serde::de::Error;
         let id: u64 = serde::de::Deserialize::deserialize(deserializer)?;
-        Ok(Entity::from_bits(id))
+        Entity::try_from_bits(id).map_err(D::Error::custom)
     }
 }
 
@@ -350,7 +380,7 @@ impl<'a> Iterator for ReserveEntitiesIterator<'a> {
             })
             .or_else(|| {
                 self.index_range.next().map(|index| Entity {
-                    generation: 0,
+                    generation: NonZeroU32::MIN,
                     index,
                 })
             })
@@ -499,7 +529,7 @@ impl Entities {
             // As `self.free_cursor` goes more and more negative, we return IDs farther
             // and farther beyond `meta.len()`.
             Entity {
-                generation: 0,
+                generation: NonZeroU32::MIN,
                 index: u32::try_from(self.meta.len() as IdCursor - n).expect("too many entities"),
             }
         }
@@ -528,7 +558,7 @@ impl Entities {
             let index = u32::try_from(self.meta.len()).expect("too many entities");
             self.meta.push(EntityMeta::EMPTY);
             Entity {
-                generation: 0,
+                generation: NonZeroU32::MIN,
                 index,
             }
         }
@@ -615,7 +645,15 @@ impl Entities {
         if meta.generation != entity.generation {
             return None;
         }
-        meta.generation += 1;
+
+        meta.generation = inc_generation_by(meta.generation, 1);
+
+        if meta.generation == NonZeroU32::MIN {
+            warn!(
+                "Entity({}) generation wrapped on Entities::free, aliasing may occur",
+                entity.index
+            );
+        }
 
         let loc = mem::replace(&mut meta.location, EntityMeta::EMPTY.location);
 
@@ -646,7 +684,7 @@ impl Entities {
     // not reallocated since the generation is incremented in `free`
     pub fn contains(&self, entity: Entity) -> bool {
         self.resolve_from_id(entity.index())
-            .map_or(false, |e| e.generation() == entity.generation)
+            .map_or(false, |e| e.generation() == entity.generation())
     }
 
     /// Clears all [`Entity`] from the World.
@@ -698,7 +736,7 @@ impl Entities {
 
         let meta = &mut self.meta[index as usize];
         if meta.location.archetype_id == ArchetypeId::INVALID {
-            meta.generation += generations;
+            meta.generation = inc_generation_by(meta.generation, generations);
             true
         } else {
             false
@@ -722,7 +760,7 @@ impl Entities {
             // Returning None handles that case correctly
             let num_pending = usize::try_from(-free_cursor).ok()?;
             (idu < self.meta.len() + num_pending).then_some(Entity {
-                generation: 0,
+                generation: NonZeroU32::MIN,
                 index,
             })
         }
@@ -840,15 +878,15 @@ impl Entities {
 #[repr(C)]
 struct EntityMeta {
     /// The current generation of the [`Entity`].
-    pub generation: u32,
+    pub generation: NonZeroU32,
     /// The current location of the [`Entity`]
     pub location: EntityLocation,
 }
 
 impl EntityMeta {
     /// meta for **pending entity**
     const EMPTY: EntityMeta = EntityMeta {
-        generation: 0,
+        generation: NonZeroU32::MIN,
         location: EntityLocation::INVALID,
     };
 }
@@ -892,14 +930,61 @@ impl EntityLocation {
     };
 }
 
+/// Offsets a generation value by the specified amount, wrapping to 1 instead of 0
+pub(crate) const fn inc_generation_by(lhs: NonZeroU32, rhs: u32) -> NonZeroU32 {
+    let (lo, hi) = lhs.get().overflowing_add(rhs);
+    let ret = lo + hi as u32;
+    // SAFETY:
+    // - Adding the overflow flag will offet overflows to start at 1 instead of 0
+    // - The sum of `NonZeroU32::MAX` + `u32::MAX` + 1 (overflow) == `NonZeroU32::MAX`
+    // - If the operation doesn't overflow, no offsetting takes place
+    unsafe { NonZeroU32::new_unchecked(ret) }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[test]
+    fn inc_generation_by_is_safe() {
+        // Correct offsets without overflow
+        assert_eq!(
+            inc_generation_by(NonZeroU32::MIN, 1),
+            NonZeroU32::new(2).unwrap()
+        );
+
+        assert_eq!(
+            inc_generation_by(NonZeroU32::MIN, 10),
+            NonZeroU32::new(11).unwrap()
+        );
+
+        // Check that overflows start at 1 correctly
+        assert_eq!(inc_generation_by(NonZeroU32::MAX, 1), NonZeroU32::MIN);
+
+        assert_eq!(
+            inc_generation_by(NonZeroU32::MAX, 2),
+            NonZeroU32::new(2).unwrap()
+        );
+
+        // Ensure we can't overflow twice
+        assert_eq!(
+            inc_generation_by(NonZeroU32::MAX, u32::MAX),
+            NonZeroU32::MAX
+        );
+    }
+
+    #[test]
+    fn entity_niche_optimization() {
+        assert_eq!(
+            std::mem::size_of::<Entity>(),
+            std::mem::size_of::<Option<Entity>>()
+        );
+    }
+
     #[test]
     fn entity_bits_roundtrip() {
         let e = Entity {
-            generation: 0xDEADBEEF,
+            generation: NonZeroU32::new(0xDEADBEEF).unwrap(),
             index: 0xBAADF00D,
         };
         assert_eq!(Entity::from_bits(e.to_bits()), e);
@@ -935,12 +1020,12 @@ mod tests {
     #[test]
     fn entity_const() {
         const C1: Entity = Entity::from_raw(42);
-        assert_eq!(42, C1.index);
-        assert_eq!(0, C1.generation);
+        assert_eq!(42, C1.index());
+        assert_eq!(1, C1.generation());
 
         const C2: Entity = Entity::from_bits(0x0000_00ff_0000_00cc);
-        assert_eq!(0x0000_00cc, C2.index);
-        assert_eq!(0x0000_00ff, C2.generation);
+        assert_eq!(0x0000_00cc, C2.index());
+        assert_eq!(0x0000_00ff, C2.generation());
 
         const C3: u32 = Entity::from_raw(33).index();
         assert_eq!(33, C3);
@@ -971,7 +1056,7 @@ mod tests {
         // The very next entity allocated should be a further generation on the same index
         let next_entity = entities.alloc();
         assert_eq!(next_entity.index(), entity.index());
-        assert!(next_entity.generation > entity.generation + GENERATIONS);
+        assert!(next_entity.generation() > entity.generation() + GENERATIONS);
     }
 
     #[test]