From 23079be23c74180bbfbaf07a69a2f79b5bcec461 Mon Sep 17 00:00:00 2001
From: Erik <somewhat.fluffy@gmail.com>
Date: Thu, 25 Mar 2021 21:06:14 +0200
Subject: [PATCH 1/2] Remove UB on overflow in `Allocate::next()`. Remove all
 unsafe in `Allocate::next()`.

---
 src/internals/entity.rs | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/src/internals/entity.rs b/src/internals/entity.rs
index c984317d..b04f0239 100644
--- a/src/internals/entity.rs
+++ b/src/internals/entity.rs
@@ -62,20 +62,27 @@ impl<'a> Iterator for Allocate {
 
     #[inline(always)]
     fn next(&mut self) -> Option<Self::Item> {
+        // Get next block whenever the block ends.
         if self.next % BLOCK_SIZE == 0 {
-            // This is either the first block, or we overflowed to the next block.
-            self.next = NEXT_ENTITY.fetch_add(BLOCK_SIZE, Ordering::Relaxed);
-            debug_assert_eq!(self.next % BLOCK_SIZE, 0);
+            static NEXT_ENTITY_BLOCK_START: AtomicU64 = AtomicU64::new(1);
+
+            // Note: since BLOCK_SIZE is divisible by 2
+            // and NEXT_ENTITY_BLOCK_START starts at 1,
+            // this skips one entity index per block.
+            // It is not the best possible performance,
+            // but zero will be skipped even on overflow and not just on start.
+            self.next = NEXT_ENTITY_BLOCK_START.fetch_add(BLOCK_SIZE, Ordering::Relaxed);
+            debug_assert_eq!(self.next % BLOCK_SIZE, 1);
         }
 
-        // Safety: self.next can't be 0 as long as the first block is skipped,
-        // and no overflow occurs in NEXT_ENTITY
-        let entity = unsafe {
-            debug_assert_ne!(self.next, 0);
-            Entity(NonZeroU64::new_unchecked(self.next))
-        };
+        // `NonZeroU64::new(self.next).map(Entity)` would give the same generated code
+        // as `unsafe { Some(unchecked) }` for *this* function.
+        // However, there may a difference when you match the resulting inlined `Option`:
+        // if its discriminant is always `Some`,
+        // then `None` case could be simply removed from generated code.
+        let entity = NonZeroU64::new(self.next).map(Entity);
         self.next += 1;
-        Some(entity)
+        entity
     }
 }
 

From 5302d4dbe2b8b2176dfd8dc783f1756ad938f50b Mon Sep 17 00:00:00 2001
From: Erik <somewhat.fluffy@gmail.com>
Date: Thu, 25 Mar 2021 21:31:54 +0200
Subject: [PATCH 2/2] Remove old unused static item.

---
 src/internals/entity.rs | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/internals/entity.rs b/src/internals/entity.rs
index b04f0239..d70a570b 100644
--- a/src/internals/entity.rs
+++ b/src/internals/entity.rs
@@ -33,10 +33,6 @@ impl Clone for Entity {
 const BLOCK_SIZE: u64 = 16;
 const BLOCK_SIZE_USIZE: usize = BLOCK_SIZE as usize;
 
-// Always divisible by BLOCK_SIZE.
-// Safety: This must never be 0, so skip the first block
-static NEXT_ENTITY: AtomicU64 = AtomicU64::new(BLOCK_SIZE);
-
 /// An iterator which yields new entity IDs.
 #[derive(Debug)]
 pub struct Allocate {