Allow Option<Entity> to leverage niche optimization

crates/bevy_ecs/src/entity/mod.rs

@@ -32,9 +32,12 @@ use crate::{archetype::ArchetypeId, storage::SparseSetIndex};
 use std::{
     convert::TryFrom,
     fmt, mem,
+    num::NonZeroU32,
     sync::atomic::{AtomicI64, Ordering},
 };
 
+use bevy_utils::tracing::warn;
+
 /// Lightweight unique ID of an entity.
 ///
 /// Obtained from [`World::spawn`](crate::world::World::spawn), typically via
@@ -46,7 +49,7 @@ use std::{
 /// [`Query::get`](crate::system::Query::get) and related methods.
 #[derive(Clone, Copy, Hash, Eq, Ord, PartialEq, PartialOrd)]
 pub struct Entity {
-    pub(crate) generation: u32,
+    pub(crate) generation: NonZeroU32,
     pub(crate) id: u32,
 }
 
@@ -57,7 +60,7 @@ pub enum AllocAtWithoutReplacement {
 }
 
 impl Entity {
-    /// Creates a new entity reference with a generation of 0.
+    /// Creates a new entity reference with a generation of 1.
     ///
     /// # Note
     ///
@@ -66,7 +69,10 @@ impl Entity {
     /// only be used for sharing entities across apps, and only when they have a scheme
     /// worked out to share an ID space (which doesn't happen by default).
     pub fn new(id: u32) -> Entity {
-        Entity { id, generation: 0 }
+        Entity {
+            id,
+            generation: GENERATION_ONE,
+        }
     }
 
     /// Convert to a form convenient for passing outside of rust.
@@ -76,17 +82,17 @@ impl Entity {
     ///
     /// No particular structure is guaranteed for the returned bits.
     pub fn to_bits(self) -> u64 {
-        u64::from(self.generation) << 32 | u64::from(self.id)
+        u64::from(self.generation()) << 32 | u64::from(self.id)
     }
 
     /// 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.
-    pub fn from_bits(bits: u64) -> Self {
-        Self {
-            generation: (bits >> 32) as u32,
+    pub fn from_bits(bits: u64) -> Option<Self> {
+        Some(Self {
+            generation: NonZeroU32::new((bits >> 32) as u32)?,
             id: bits as u32,
-        }
+        })
     }
 
     /// Return a transiently unique identifier.
@@ -104,7 +110,7 @@ impl Entity {
     /// given id has been reused (id, generation) pairs uniquely identify a given Entity.
     #[inline]
     pub fn generation(self) -> u32 {
-        self.generation
+        self.generation.get()
     }
 }
 
@@ -144,10 +150,15 @@ impl<'a> Iterator for ReserveEntitiesIterator<'a> {
         self.id_iter
             .next()
             .map(|&id| Entity {
-                generation: self.meta[id as usize].generation,
+                generation: self.meta[id as usize].generation.unwrap(),
                 id,
             })
-            .or_else(|| self.id_range.next().map(|id| Entity { generation: 0, id }))
+            .or_else(|| {
+                self.id_range.next().map(|id| Entity {
+                    generation: GENERATION_ONE,
+                    id,
+                })
+            })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -255,7 +266,7 @@ impl Entities {
             // Allocate from the freelist.
             let id = self.pending[(n - 1) as usize];
             Entity {
-                generation: self.meta[id as usize].generation,
+                generation: self.meta[id as usize].generation.unwrap(), // Safe, meta from pending list, so generation is valid
                 id,
             }
         } else {
@@ -265,7 +276,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: GENERATION_ONE,
                 id: u32::try_from(self.meta.len() as i64 - n).expect("too many entities"),
             }
         }
@@ -287,13 +298,16 @@ impl Entities {
             let new_free_cursor = self.pending.len() as i64;
             *self.free_cursor.get_mut() = new_free_cursor;
             Entity {
-                generation: self.meta[id as usize].generation,
+                generation: self.meta[id as usize].generation.unwrap(), // Safe, meta from pending list, so generation is valid
                 id,
             }
         } else {
             let id = u32::try_from(self.meta.len()).expect("too many entities");
             self.meta.push(EntityMeta::EMPTY);
-            Entity { generation: 0, id }
+            Entity {
+                generation: GENERATION_ONE,
+                id,
+            }
         }
     }
 
@@ -324,7 +338,7 @@ impl Entities {
             ))
         };
 
-        self.meta[entity.id as usize].generation = entity.generation;
+        self.meta[entity.id as usize].generation = Some(entity.generation);
 
         loc
     }
@@ -352,14 +366,14 @@ impl Entities {
             let current_meta = &mut self.meta[entity.id as usize];
             if current_meta.location.archetype_id == ArchetypeId::INVALID {
                 AllocAtWithoutReplacement::DidNotExist
-            } else if current_meta.generation == entity.generation {
+            } else if current_meta.generation == Some(entity.generation) {
                 AllocAtWithoutReplacement::Exists(current_meta.location)
             } else {
                 return AllocAtWithoutReplacement::ExistsWithWrongGeneration;
             }
         };
 
-        self.meta[entity.id as usize].generation = entity.generation;
+        self.meta[entity.id as usize].generation = Some(entity.generation);
         result
     }
 
@@ -370,19 +384,26 @@ impl Entities {
         self.verify_flushed();
 
         let meta = &mut self.meta[entity.id as usize];
-        if meta.generation != entity.generation {
+        if meta.generation != Some(entity.generation) {
             return None;
         }
-        meta.generation += 1;
 
-        let loc = mem::replace(&mut meta.location, EntityMeta::EMPTY.location);
+        meta.generation = NonZeroU32::new(meta.generation.unwrap().get().wrapping_add(1));
 
-        self.pending.push(entity.id);
+        if meta.generation.is_some() {
+            self.pending.push(entity.id);
 
-        let new_free_cursor = self.pending.len() as i64;
-        *self.free_cursor.get_mut() = new_free_cursor;
-        self.len -= 1;
-        Some(loc)
+            let new_free_cursor = self.pending.len() as i64;
+            *self.free_cursor.get_mut() = new_free_cursor;
+            self.len -= 1;
+        } else {
+            warn!(
+                "Entity generation exhausted. Retiring slot for entity #{}",
+                entity.id
+            );
+        }
+
+        Some(mem::replace(&mut meta.location, EntityMeta::EMPTY.location))
     }
 
     /// Ensure at least `n` allocations can succeed without reallocating.
@@ -397,11 +418,11 @@ impl Entities {
     }
 
     /// Returns true if the [`Entities`] contains [`entity`](Entity).
-    // This will return false for entities which have been freed, even if
-    // not reallocated since the generation is incremented in `free`
+    /// This will return false for entities which have been freed, even if
+    /// not reallocated since the generation is incremented in `free`
     pub fn contains(&self, entity: Entity) -> bool {
         self.resolve_from_id(entity.id())
-            .map_or(false, |e| e.generation() == entity.generation)
+            .map_or(false, |e| e.generation == entity.generation)
     }
 
     pub fn clear(&mut self) {
@@ -415,7 +436,8 @@ impl Entities {
     pub fn get(&self, entity: Entity) -> Option<EntityLocation> {
         if (entity.id as usize) < self.meta.len() {
             let meta = &self.meta[entity.id as usize];
-            if meta.generation != entity.generation
+            if meta.generation.is_none()
+                || meta.generation.unwrap() != entity.generation
                 || meta.location.archetype_id == ArchetypeId::INVALID
             {
                 return None;
@@ -435,14 +457,17 @@ impl Entities {
     pub fn resolve_from_id(&self, id: u32) -> Option<Entity> {
         let idu = id as usize;
         if let Some(&EntityMeta { generation, .. }) = self.meta.get(idu) {
-            Some(Entity { generation, id })
+            generation.map(|generation| Entity { generation, id })
         } else {
             // `id` is outside of the meta list - check whether it is reserved but not yet flushed.
             let free_cursor = self.free_cursor.load(Ordering::Relaxed);
             // If this entity was manually created, then free_cursor might be positive
             // Returning None handles that case correctly
             let num_pending = usize::try_from(-free_cursor).ok()?;
-            (idu < self.meta.len() + num_pending).then(|| Entity { generation: 0, id })
+            (idu < self.meta.len() + num_pending).then(|| Entity {
+                generation: GENERATION_ONE,
+                id,
+            })
         }
     }
 
@@ -469,13 +494,15 @@ impl Entities {
             self.meta.resize(new_meta_len, EntityMeta::EMPTY);
             self.len += -current_free_cursor as u32;
             for (id, meta) in self.meta.iter_mut().enumerate().skip(old_meta_len) {
-                init(
-                    Entity {
-                        id: id as u32,
-                        generation: meta.generation,
-                    },
-                    &mut meta.location,
-                );
+                if let Some(generation) = meta.generation {
+                    init(
+                        Entity {
+                            id: id as u32,
+                            generation,
+                        },
+                        &mut meta.location,
+                    );
+                }
             }
 
             *free_cursor = 0;
@@ -488,7 +515,7 @@ impl Entities {
             init(
                 Entity {
                     id,
-                    generation: meta.generation,
+                    generation: meta.generation.unwrap(), // Safe, meta from pending list, so generation is valid
                 },
                 &mut meta.location,
             );
@@ -518,13 +545,13 @@ impl Entities {
 
 #[derive(Copy, Clone, Debug)]
 pub struct EntityMeta {
-    pub generation: u32,
+    pub generation: Option<NonZeroU32>,
     pub location: EntityLocation,
 }
 
 impl EntityMeta {
     const EMPTY: EntityMeta = EntityMeta {
-        generation: 0,
+        generation: Some(GENERATION_ONE),
         location: EntityLocation {
             archetype_id: ArchetypeId::INVALID,
             index: usize::MAX, // dummy value, to be filled in
@@ -542,17 +569,67 @@ pub struct EntityLocation {
     pub index: usize,
 }
 
+// Constant for the initial generation value, removes the need to unwrap in opt-0 code but doesn't
+// require unsafe. When rust 1.57 drops, we can exchange [][1] for panic!().
+const GENERATION_ONE: NonZeroU32 = if let Some(gen) = NonZeroU32::new(1) {
+    gen
+} else {
+    [][1]
+};
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn entity_bits_roundtrip() {
         let e = Entity {
-            generation: 0xDEADBEEF,
+            generation: NonZeroU32::new(0xDEADBEEF).unwrap(),
             id: 0xBAADF00D,
         };
-        assert_eq!(Entity::from_bits(e.to_bits()), e);
+        assert_eq!(Entity::from_bits(e.to_bits()).unwrap(), e);
+    }
+
+    #[test]
+    fn entity_bad_bits() {
+        let bits: u64 = 0xBAADF00D;
+        assert_eq!(Entity::from_bits(bits), None);
+    }
+
+    #[test]
+    fn entity_option_size_optimized() {
+        assert_eq!(
+            core::mem::size_of::<Option<Entity>>(),
+            core::mem::size_of::<Entity>()
+        );
+    }
+
+    #[test]
+    fn entities_generation_increment() {
+        let mut entities = Entities::default();
+
+        let entity_old = entities.alloc();
+        entities.free(entity_old);
+        let entity_new = entities.alloc();
+
+        assert_eq!(entity_old.id, entity_new.id);
+        assert_eq!(entity_old.generation() + 1, entity_new.generation());
+    }
+
+    #[test]
+    fn entities_generation_overflow() {
+        let mut entities = Entities::default();
+        let mut entity_old = entities.alloc();
+
+        // Modify generation on entity and entities to cause overflow on free
+        entity_old.generation = NonZeroU32::new(u32::MAX).unwrap();
+        entities.meta[entity_old.id as usize].generation = Some(entity_old.generation);
+        entities.free(entity_old);
+
+        // Request new entity, we shouldn't get the one we just freed since it overflowed
+        let entity_new = entities.alloc();
+        assert!(entity_old.id != entity_new.id);
+        assert_eq!(entity_new.generation(), 1);
     }
 
     #[test]