From 207a3f7861200da75bf833eb59585dd8d88f74e0 Mon Sep 17 00:00:00 2001
From: marc0246 <40955683+marc0246@users.noreply.github.com>
Date: Sun, 20 Oct 2024 13:41:18 +0200
Subject: [PATCH] Use `NonZeroUsize` for `slab_capacity`

---
 src/lib.rs | 35 +++++++++++++++++++----------------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0a13e4b..60e61e7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -32,6 +32,7 @@ use core::{
     cell::Cell,
     fmt,
     mem::{self, ManuallyDrop, MaybeUninit},
+    num::NonZeroUsize,
     ptr::{self, NonNull},
 };
 
@@ -140,7 +141,7 @@ pub struct SlabAllocator<T> {
     /// Points to the end of the slab of the head of the slab list. If the slab list is empty then
     /// this is dangling.
     free_end: Cell<NonNull<Slot<T>>>,
-    slab_capacity: usize,
+    slab_capacity: NonZeroUsize,
 }
 
 // SAFETY: The pointers we hold are not referencing the stack or TLS or anything like that, they
@@ -168,16 +169,18 @@ impl<T> SlabAllocator<T> {
     #[inline]
     #[must_use]
     pub const fn new(slab_capacity: usize) -> Self {
-        assert!(slab_capacity != 0);
-
-        let dangling = NonNull::dangling();
-
-        SlabAllocator {
-            free_list_head: Cell::new(None),
-            slab_list_head: Cell::new(None),
-            free_start: Cell::new(dangling),
-            free_end: Cell::new(dangling),
-            slab_capacity,
+        if let Some(slab_capacity) = NonZeroUsize::new(slab_capacity) {
+            let dangling = NonNull::dangling();
+
+            SlabAllocator {
+                free_list_head: Cell::new(None),
+                slab_list_head: Cell::new(None),
+                free_start: Cell::new(dangling),
+                free_end: Cell::new(dangling),
+                slab_capacity,
+            }
+        } else {
+            panic!("`slab_capacity` must be non-zero");
         }
     }
 
@@ -297,8 +300,7 @@ impl<T> SlabAllocator<T> {
 
         // SAFETY:
         // * We know that the offset must be in bounds of the allocated object because we allocated
-        //   `self.slab_capacity` slots, and by our own invariant, `self.slab_capacity` must be
-        //   non-zero.
+        //   `self.slab_capacity` slots and `self.slab_capacity` is non-zero.
         // * The computed offset cannot overflow an `isize` because we used `Layout` for the layout
         //   calculation.
         // * The computed offset cannot wrap around the address space for the same reason as the
@@ -306,7 +308,8 @@ impl<T> SlabAllocator<T> {
         let free_start = unsafe { NonNull::new_unchecked(slots.as_ptr().add(1)) };
 
         // SAFETY: Same as the previous.
-        let free_end = unsafe { NonNull::new_unchecked(slots.as_ptr().add(self.slab_capacity)) };
+        let free_end =
+            unsafe { NonNull::new_unchecked(slots.as_ptr().add(self.slab_capacity.get())) };
 
         self.free_start.set(free_start);
         self.free_end.set(free_end);
@@ -333,7 +336,7 @@ impl<T> SlabAllocator<T> {
 
     fn slab_layout(&self) -> Layout {
         Layout::new::<Option<NonNull<Slab<T>>>>()
-            .extend(Layout::array::<Slot<T>>(self.slab_capacity).unwrap())
+            .extend(Layout::array::<Slot<T>>(self.slab_capacity.get()).unwrap())
             .unwrap()
             .0
             .pad_to_align()
@@ -383,7 +386,7 @@ impl<T> SlabAllocator<T> {
                 //   the layout calculation.
                 // * The computed offset cannot wrap around the address space for the same reason
                 //   as the previous.
-                let slots_end = unsafe { slots_start.add(self.slab_capacity) };
+                let slots_end = unsafe { slots_start.add(self.slab_capacity.get()) };
 
                 if (slots_start..slots_end).contains(&ptr.as_ptr()) {
                     break;