foundation: slab allocator

crates/foundation/src/alloc/slab.rs

@@ -1,8 +1,8 @@
 //! A slab allocator
 
-use core::{alloc::Layout, cell::Cell, hint::unlikely, mem::offset_of, num::NonZero, ptr::NonNull};
-
-const UNLINKED: NonNull<()> = unsafe { NonNull::new_unchecked(!0 as *mut ()) };
+use core::{
+    alloc::Layout, cell::Cell, fmt::Debug, hint::unlikely, num::NonZero, pin::Pin, ptr::NonNull,
+};
 
 use liballoc::alloc::Allocator;
 
@@ -15,7 +15,20 @@ pub struct Slab<A: Allocator + Clone> {
     element_size: usize,
     /// Pointer to the first chunk in the slab.
     head: Option<NonNull<SlabChunk<A>>>,
+    /// Pointer to the first chunk in the slab that is full.
+    full_head: Option<NonNull<SlabChunk<A>>>,
     alloc: A,
+    _pd: core::marker::PhantomPinned,
+}
+
+impl<A: Allocator + Clone> Debug for Slab<A> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("Slab")
+            .field("element_size", &self.element_size)
+            .field("head", &self.head)
+            .field("full_head", &self.full_head)
+            .finish()
+    }
 }
 
 struct SlabChunk<A: Allocator + Clone> {
@@ -39,7 +52,18 @@ enum SlotResult {
     None,
 }
 
+enum FreeSlotResult {
+    Empty,
+    WasFull,
+    NotEmpty,
+}
+
 impl<A: Allocator + Clone> SlabChunk<A> {
+    unsafe fn slab_pinned_mut<'a>(&self) -> Pin<&'a mut Slab<A>> {
+        // SAFETY: `slab` is created from a `Pin<&mut Slab<A>>`
+        unsafe { Pin::new_unchecked(self.slab.as_ptr().as_mut_unchecked()) }
+    }
+
     fn pop_free_slot(&self) -> SlotResult {
         let Some(slot) = self.free.get() else {
             return SlotResult::None;
@@ -55,14 +79,18 @@ impl<A: Allocator + Clone> SlabChunk<A> {
     }
 
     /// returns `true` if the slab is now empty and can be freed
-    fn push_free_slot(&self, slot: NonNull<u8>) -> bool {
+    fn push_free_slot(&self, slot: NonNull<u8>) -> FreeSlotResult {
         let slot = slot.cast::<ChunkSlot>();
         let next = self.free.get();
         unsafe { slot.as_ptr().write(ChunkSlot(next)) };
         self.free.set(Some(slot));
         self.count.update(|count| count - 1);
 
-        self.count.get() == 0
+        match (next, self.count.get()) {
+            (None, _) => FreeSlotResult::WasFull,
+            (_, 0) => FreeSlotResult::Empty,
+            _ => FreeSlotResult::NotEmpty,
+        }
     }
 }
 
@@ -79,57 +107,72 @@ impl<A: Allocator + Clone> Slab<A> {
             "element_size must be a power of two"
         );
 
-        Self {
+        let slab = Self {
             element_size,
             head: None,
-            alloc,
-        }
+            full_head: None,
+            alloc: alloc.clone(),
+            _pd: core::marker::PhantomPinned,
+        };
+
+        #[cfg(test)]
+        std::eprintln!("Slab::new({element_size}) -> {slab:#?}",);
+
+        slab
     }
 
-    fn first_slot_offset(&self) -> usize {
-        mem::align_up(core::mem::size_of::<SlabChunk<A>>(), self.element_size)
+    fn first_slot_offset(element_size: usize) -> usize {
+        mem::align_up(core::mem::size_of::<SlabChunk<A>>(), element_size)
     }
 
-    fn count_and_layout(&self) -> (usize, Layout) {
+    fn count_and_layout(element_size: usize) -> (usize, Layout) {
         let (count, size, align) = {
-            let one_page_count = (PAGE_SIZE - self.first_slot_offset()) / self.element_size;
+            let one_page_count = (PAGE_SIZE - Self::first_slot_offset(element_size)) / element_size;
 
             if one_page_count < 3 {
                 let count = 3;
-                let size =
-                    (self.first_slot_offset() + count * self.element_size).next_power_of_two();
+                let size = (Self::first_slot_offset(element_size) + count * element_size)
+                    .next_power_of_two();
 
                 assert!(size.is_multiple_of(PAGE_SIZE));
                 assert!(size >= PAGE_SIZE);
-                assert!(size.is_multiple_of(self.element_size));
+                //assert!(size.is_multiple_of(self.element_size));
 
                 (count, size, size)
             } else {
                 let count = one_page_count;
                 let size = PAGE_SIZE;
-                (count, size, self.element_size)
+                let align = element_size.max(PAGE_SIZE);
+                (count, size, align)
             }
         };
 
+        #[cfg(test)]
+        std::eprintln!("element_size: {element_size} count: {count}, size: {size}, align: {align}");
+
         (count, unsafe {
             Layout::from_size_align_unchecked(size, align)
         })
     }
 
-    fn alloc_chunk(&mut self) -> NonNull<SlabChunk<A>> {
+    fn alloc_chunk(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
         // we want to limit chunks to 1 page unless the element size is so
         // large that we can fit fewer than 3 elements in a page.
-        let (count, layout) = self.count_and_layout();
+        let (count, layout) = Self::count_and_layout(self.element_size);
 
         let Some(bytes) = self.alloc.allocate(layout).ok() else {
             panic!()
         };
 
         let chunk = bytes.as_non_null_ptr().cast::<SlabChunk<A>>();
+
+        #[cfg(test)]
+        std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})");
+
         unsafe {
             let first_slot = chunk
                 .as_ptr()
-                .byte_add(self.first_slot_offset())
+                .byte_add(Self::first_slot_offset(self.element_size))
                 .cast::<ChunkSlot>();
 
             for i in 0..(count - 1) {
@@ -143,7 +186,8 @@ impl<A: Allocator + Clone> Slab<A> {
 
             chunk.write(SlabChunk {
                 next: self.head,
-                slab: NonNull::from(self),
+                // SAFETY: we only access `slab` via `SlabChunk::slab_pinned_mut`
+                slab: NonNull::from(Pin::into_inner_unchecked(self)),
                 free: Cell::new(Some(NonNull::new_unchecked(first_slot))),
                 count: Cell::new(0),
             });
@@ -153,26 +197,37 @@ impl<A: Allocator + Clone> Slab<A> {
     }
 
     #[cold]
-    fn alloc_chunk_cold(&mut self) -> NonNull<SlabChunk<A>> {
+    fn alloc_chunk_cold(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
         self.alloc_chunk()
     }
 
-    fn alloc_slot(&mut self) -> NonNull<[u8]> {
-        let mut chunk = match self.head {
+    fn alloc_slot(mut self: Pin<&mut Self>) -> NonNull<[u8]> {
+        #[cfg(test)]
+        std::eprintln!("Slab::alloc_slot({self:?})");
+
+        let chunk = match self.head {
             Some(chunk) => chunk,
             None => {
-                let chunk = self.alloc_chunk_cold();
-                self.head = Some(chunk);
+                let chunk = self.as_mut().alloc_chunk_cold();
+                unsafe {
+                    self.as_mut().get_unchecked_mut().head = Some(chunk);
+                }
                 chunk
             }
         };
 
-        let chunk = unsafe { chunk.as_mut() };
+        let chunk = unsafe { chunk.as_ptr().as_mut_unchecked() };
 
         let ptr = match chunk.pop_free_slot() {
             SlotResult::Some(non_null) => non_null,
             SlotResult::Last(non_null) => {
-                self.head = chunk.next.replace(UNLINKED.cast());
+                unsafe {
+                    let mut_ref = self.as_mut().get_unchecked_mut();
+                    mut_ref.head = chunk.next.take();
+                    chunk.next = mut_ref.full_head;
+                    mut_ref.full_head = Some(chunk.into());
+                }
+
                 non_null
             }
             SlotResult::None => {
@@ -183,49 +238,97 @@ impl<A: Allocator + Clone> Slab<A> {
         ptr.cast_slice(self.element_size)
     }
 
-    fn free_slot(&mut self, slot: NonNull<u8>) {
-        let (_, layout) = self.count_and_layout();
-        let mut chunk = slot
+    fn free_slot(element_size: usize, slot: NonNull<u8>) {
+        let (_, layout) = Self::count_and_layout(element_size);
+
+        let chunk_ptr = slot
             .map_addr(|addr| unsafe {
                 NonZero::new_unchecked(mem::align_down(addr.get(), layout.align()))
             })
             .cast::<SlabChunk<A>>();
 
-        let chunk = unsafe { chunk.as_mut() };
-        let linked = chunk.next != Some(UNLINKED.cast());
+        #[cfg(test)]
+        std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})");
 
-        if chunk.push_free_slot(slot) {
-            if linked {
-                let mut head = self.head.expect("chunk is linked, so head exists");
-                while let Some(next) = unsafe { head.as_ref().next } {
+        let chunk = unsafe { chunk_ptr.as_ptr().as_mut_unchecked() };
+
+        let mut slab = unsafe { chunk.slab_pinned_mut() };
+
+        match chunk.push_free_slot(slot) {
+            FreeSlotResult::Empty => {
+                // chunk is empty:
+                // first unlink the chunk from the slab..
+                if slab.head == Some(chunk.into()) {
+                    unsafe { slab.as_mut().get_unchecked_mut().head = chunk.next };
+                } else {
+                    let mut head = slab.head.expect("chunk is linked, so head exists");
+                    while let Some(next) = unsafe { (&raw const (*head.as_ptr()).next).read() } {
                         if next == chunk.into() {
-                        unsafe { head.as_mut().next = chunk.next };
+                            unsafe { (&raw mut (*head.as_ptr()).next).write(chunk.next) };
                             break;
                         }
                         head = next;
                     }
                 }
 
+                // ..then free it.
+                #[cfg(test)]
+                std::eprintln!("free_slot::drop({chunk_ptr:?}, layout: {layout:?})");
+
+                unsafe { slab.alloc.deallocate(chunk_ptr.cast(), layout) };
+            }
+            FreeSlotResult::WasFull => {
+                // chunk was full:
+
+                // unlink it from the full list
+                if slab.full_head == Some(chunk.into()) {
+                    unsafe { slab.as_mut().get_unchecked_mut().full_head = chunk.next };
+                } else {
+                    let mut head = slab.full_head.expect("chunk is linked, so head exists");
+                    while let Some(next) = unsafe { (&raw const (*head.as_ptr()).next).read() } {
+                        if next == chunk.into() {
+                            unsafe { (&raw mut (*head.as_ptr()).next).write(chunk.next) };
+                            break;
+                        }
+                        head = next;
+                    }
+                }
+
+                // link it to the head of the slab
+                chunk.next = slab.head;
                 unsafe {
-                self.alloc
-                    .deallocate(NonNull::from_mut(chunk).cast(), layout)
-            };
-        } else if !linked {
-            chunk.next = self.head;
-            self.head = Some(chunk.into());
+                    slab.as_mut().get_unchecked_mut().head = Some(chunk.into());
+                }
+            }
+            FreeSlotResult::NotEmpty => {
+                // chunk is not empty, and was not full, so nothing to do
+            }
         }
     }
 }
 
 impl<A: Allocator + Clone> Drop for Slab<A> {
     fn drop(&mut self) {
-        let (_, layout) = self.count_and_layout();
+        let (_, layout) = Self::count_and_layout(self.element_size);
 
         let mut chunk = self.head.take();
         while let Some(mut chunk_ptr) = chunk {
             let chunk_ref = unsafe { chunk_ptr.as_mut() };
             chunk = chunk_ref.next.take();
 
+            #[cfg(test)]
+            std::eprintln!("Slab::drop({chunk_ptr:?}, layout: {layout:?})");
+
+            unsafe {
+                self.alloc.deallocate(chunk_ptr.cast(), layout);
+            }
+        }
+
+        let mut chunk = self.full_head.take();
+        while let Some(mut chunk_ptr) = chunk {
+            let chunk_ref = unsafe { chunk_ptr.as_mut() };
+            chunk = chunk_ref.next.take();
+
             unsafe {
                 self.alloc.deallocate(chunk_ptr.cast(), layout);
             }
@@ -242,7 +345,8 @@ pub struct SlabAllocator<A: Allocator + Clone> {
 
 impl<A: Allocator + Clone> SlabAllocator<A> {
     pub fn new(alloc: A) -> Self {
-        let slabs = [
+        Self {
+            slabs: [
                 Slab::new(16, alloc.clone()),
                 Slab::new(32, alloc.clone()),
                 Slab::new(64, alloc.clone()),
@@ -251,9 +355,9 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
                 Slab::new(512, alloc.clone()),
                 Slab::new(1024, alloc.clone()),
                 Slab::new(2048, alloc.clone()),
-        ];
-
-        Self { slabs, alloc }
+            ],
+            alloc,
+        }
     }
 
     fn slab_index_for_size(size: usize) -> Option<usize> {
@@ -268,6 +372,14 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
         // subtracting 4 (since 2^4 = 16)
         let index = (size.next_power_of_two().trailing_zeros() - 16usize.trailing_zeros()) as usize;
 
+        #[cfg(test)]
+        std::eprintln!(
+            "slab_index_for_size: size: {size}, next_power_of_two: {}, trailing_zeros: {}, index: {}",
+            size.next_power_of_two(),
+            size.next_power_of_two().trailing_zeros(),
+            index,
+        );
+
         // we have 8 slabs, so the index must be less than 8
         if index < SLAB_ALLOCATOR_BUCKETS {
             Some(index)
@@ -276,7 +388,10 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
         }
     }
 
-    pub fn alloc(&mut self, layout: Layout) -> Option<NonNull<[u8]>> {
+    pub fn alloc(self: Pin<&mut Self>, layout: Layout) -> Option<NonNull<[u8]>> {
+        #[cfg(test)]
+        std::eprintln!("SlabAllocator::alloc({:?})", layout);
+
         if unlikely(layout.size() == 0) {
             return Some(NonNull::dangling().cast_slice(0));
         }
@@ -284,9 +399,12 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
         let size = layout.size().max(layout.align());
 
         match Self::slab_index_for_size(size) {
-            Some(slab_index) => {
-                Some(unsafe { self.slabs.get_unchecked_mut(slab_index).alloc_slot() })
-            }
+            Some(slab_index) => unsafe {
+                Some(
+                    self.map_unchecked_mut(|alloc| alloc.slabs.get_unchecked_mut(slab_index))
+                        .alloc_slot(),
+                )
+            },
             None => {
                 // allocate directly from the backing allocator
                 self.alloc.allocate(layout).ok()
@@ -294,7 +412,10 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
         }
     }
 
-    pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
+    pub fn dealloc(self: Pin<&mut Self>, ptr: NonNull<u8>, layout: Layout) {
+        #[cfg(test)]
+        std::eprintln!("SlabAllocator::dealloc({:?})", layout);
+
         if unlikely(layout.size() == 0) {
             return;
         }
@@ -302,15 +423,8 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
         let size = layout.size().max(layout.align());
 
         match Self::slab_index_for_size(size) {
-            Some(_) => unsafe {
-                let chunk = ptr
-                    .map_addr(|addr| {
-                        NonZero::new_unchecked(mem::align_down(addr.get(), layout.align()))
-                    })
-                    .cast::<SlabChunk<A>>();
-
-                let mut slab = (&raw const (*chunk.as_ptr()).slab).read();
-                slab.as_mut().free_slot(ptr);
+            Some(slab_index) => unsafe {
+                Slab::<A>::free_slot(self.slabs.get_unchecked(slab_index).element_size, ptr);
             },
             None => {
                 // deallocate directly to the backing allocator
@@ -320,30 +434,16 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
     }
 }
 
-impl<A: Allocator + Clone> Drop for SlabAllocator<A> {
-    fn drop(&mut self) {}
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    use core::{mem::forget, pin};
     use std::{prelude::rust_2024::*, sync::Mutex};
 
-    struct SlabAllocatorWrapper<A: Allocator + Clone>(Mutex<SlabAllocator<A>>);
+    struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex<Pin<&'a mut SlabAllocator<A>>>);
 
-    impl<A: Allocator + Clone> SlabAllocatorWrapper<A> {
-        fn new(alloc: A) -> Self {
-            Self(Mutex::new(SlabAllocator::new(alloc)))
-        }
-    }
-    impl SlabAllocatorWrapper<std::alloc::System> {
-        fn new_system() -> SlabAllocatorWrapper<std::alloc::System> {
-            SlabAllocatorWrapper::new(std::alloc::System)
-        }
-    }
-
-    unsafe impl<A: Allocator + Clone> Allocator for SlabAllocatorWrapper<A> {
+    unsafe impl<'a, A: Allocator + Clone> Allocator for SlabAllocatorWrapper<'a, A> {
         fn allocate(
             &self,
             layout: core::alloc::Layout,
@@ -351,29 +451,33 @@ mod tests {
             self.0
                 .lock()
                 .unwrap()
+                .as_mut()
                 .alloc(layout)
                 .ok_or(core::alloc::AllocError)
         }
 
         unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: core::alloc::Layout) {
-            self.0.lock().unwrap().dealloc(ptr, layout)
+            self.0.lock().unwrap().as_mut().dealloc(ptr, layout)
         }
     }
 
     #[test]
     fn test_slab_allocator() {
-        let alloc = SlabAllocatorWrapper::new_system();
+        let mut slab = pin::pin!(SlabAllocator::new(std::alloc::Global));
+        let alloc = SlabAllocatorWrapper(Mutex::new(slab.as_mut()));
 
-        let layout = Layout::from_size_align(16, 8).unwrap();
+        let layout = Layout::from_size_align(4, 4).unwrap();
         let ptr = alloc.allocate(layout).unwrap();
         assert_eq!(ptr.len(), 16);
         assert_eq!(ptr.addr().get() % 8, 0);
 
+        {
             let mut boxed = Box::new_in(42u32, &alloc);
             assert_eq!(*boxed, 42);
             assert_ne!(Box::as_non_null(&mut boxed), ptr.cast());
-
             unsafe { alloc.deallocate(ptr.cast(), layout) };
-        _ = boxed;
+            drop(boxed);
+            std::eprintln!("test_slab_allocator: dropped boxed");
+        }
     }
 }