From ef7f9241432302b399f68ac474907b3e68400a6c Mon Sep 17 00:00:00 2001 From: janis Date: Sat, 8 Aug 2026 22:05:40 +0200 Subject: [PATCH] slab alloc: use linked list --- crates/foundation/src/alloc/slab.rs | 172 ++++++++++++++-------------- 1 file changed, 83 insertions(+), 89 deletions(-) diff --git a/crates/foundation/src/alloc/slab.rs b/crates/foundation/src/alloc/slab.rs index b718591..556962e 100644 --- a/crates/foundation/src/alloc/slab.rs +++ b/crates/foundation/src/alloc/slab.rs @@ -6,7 +6,10 @@ use core::{ use liballoc::alloc::Allocator; -use crate::mem; +use crate::{ + collections::linked_list::{LinkedList, LinkedListNode}, + mem, +}; const PAGE_SIZE: usize = 4096; @@ -14,9 +17,9 @@ pub struct Slab { /// Size and alignment of each element in the slab. element_size: usize, /// Pointer to the first chunk in the slab. - head: Option>>, + chunks: LinkedList>, /// Pointer to the first chunk in the slab that is full. - full_head: Option>>, + full_chunks: LinkedList>, alloc: A, _pd: core::marker::PhantomPinned, } @@ -25,12 +28,11 @@ impl Debug for Slab { 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() + .finish_non_exhaustive() } } +#[derive(Debug)] struct SlabChunk { /// Pointer to the next chunk in the slab. next: Option>>, @@ -44,6 +46,16 @@ struct SlabChunk { count: Cell, } +unsafe impl LinkedListNode for SlabChunk { + fn next(this: NonNull) -> Option> { + unsafe { (&raw const (*this.as_ptr()).next).read() } + } + + fn set_next(this: NonNull, next: Option>) { + unsafe { (&raw mut (*this.as_ptr()).next).write(next) }; + } +} + struct ChunkSlot(Option>); enum SlotResult { @@ -109,8 +121,8 @@ impl Slab { let slab = Self { element_size, - head: None, - full_head: None, + chunks: LinkedList::new(), + full_chunks: LinkedList::new(), alloc: alloc.clone(), _pd: core::marker::PhantomPinned, }; @@ -155,16 +167,17 @@ impl Slab { }) } - fn alloc_chunk(self: Pin<&mut Self>) -> NonNull> { + fn alloc_chunk(element_size: usize, alloc: &A, slab: NonNull) -> NonNull> { // 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(self.element_size); + let (count, layout) = Self::count_and_layout(element_size); - let Some(bytes) = self.alloc.allocate(layout).ok() else { + let Some(bytes) = alloc.allocate(layout).ok() else { panic!() }; let chunk = bytes.as_non_null_ptr().cast::>(); + let _ = chunk.as_ptr().expose_provenance(); #[cfg(test)] std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})"); @@ -172,22 +185,22 @@ impl Slab { unsafe { let first_slot = chunk .as_ptr() - .byte_add(Self::first_slot_offset(self.element_size)) + .byte_add(Self::first_slot_offset(element_size)) .cast::(); for i in 0..(count - 1) { - let chunk = first_slot.byte_add(i * self.element_size); - let next = first_slot.byte_add((i + 1) * self.element_size); + let chunk = first_slot.byte_add(i * element_size); + let next = first_slot.byte_add((i + 1) * element_size); chunk.write(ChunkSlot(Some(NonNull::new_unchecked(next)))); } first_slot - .byte_add((count - 1) * self.element_size) + .byte_add((count - 1) * element_size) .write(ChunkSlot(None)); chunk.write(SlabChunk { - next: self.head, + next: None, // SAFETY: we only access `slab` via `SlabChunk::slab_pinned_mut` - slab: NonNull::from(Pin::into_inner_unchecked(self)), + slab, free: Cell::new(Some(NonNull::new_unchecked(first_slot))), count: Cell::new(0), }); @@ -196,36 +209,31 @@ impl Slab { chunk } - #[cold] - fn alloc_chunk_cold(self: Pin<&mut Self>) -> NonNull> { - self.alloc_chunk() - } - 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.as_mut().alloc_chunk_cold(); - unsafe { - self.as_mut().get_unchecked_mut().head = Some(chunk); - } - chunk - } + let slab: NonNull = unsafe { self.as_mut().get_unchecked_mut().into() }; + let alloc = self.alloc.clone(); + let element_size = self.element_size; + + let mut chunk = unsafe { + self.as_mut() + .get_unchecked_mut() + .chunks + .get_or_insert_front_with(move || Self::alloc_chunk(element_size, &alloc, slab)) }; - let chunk = unsafe { chunk.as_ptr().as_mut_unchecked() }; + let chunk_ref = chunk.get_mut(); - let ptr = match chunk.pop_free_slot() { + let ptr = match chunk_ref.pop_free_slot() { SlotResult::Some(non_null) => non_null, SlotResult::Last(non_null) => { + let chunk = chunk.remove(); 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()); + self.as_mut() + .map_unchecked_mut(|slab| &mut slab.full_chunks) + .push_front(chunk); } non_null @@ -241,11 +249,14 @@ impl Slab { fn free_slot(element_size: usize, slot: NonNull) { 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::>(); + let addr = + unsafe { NonZero::new_unchecked(mem::align_down(slot.addr().get(), layout.align())) }; + + // grab provenance exposed in alloc_chunk, since `slot` may be noalias restricted to its layout. + + // let chunk_ptr: NonNull> = NonNull::with_exposed_provenance(addr); + + let chunk_ptr: NonNull> = slot.with_addr(addr).cast(); #[cfg(test)] std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})"); @@ -258,47 +269,35 @@ impl Slab { 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 { (&raw mut (*head.as_ptr()).next).write(chunk.next) }; - break; - } - head = next; - } + if let Some(chunk) = + unsafe { slab.as_mut().map_unchecked_mut(|slab| &mut slab.chunks) } + .remove_if(|c| c == chunk_ptr) + .next() + { + // ..then free it. + #[cfg(test)] + std::eprintln!("free_slot::drop({chunk:?}, layout: {layout:?})"); + + unsafe { slab.alloc.deallocate(chunk.cast(), layout) }; } - - // ..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 { - slab.as_mut().get_unchecked_mut().head = Some(chunk.into()); + if let Some(chunk) = unsafe { + slab.as_mut() + .map_unchecked_mut(|slab| &mut slab.full_chunks) } + .remove_if(|c| c == chunk_ptr) + .next() + { + unsafe { + slab.as_mut() + .map_unchecked_mut(|slab| &mut slab.chunks) + .push_front(chunk) + }; + }; } FreeSlotResult::NotEmpty => { // chunk is not empty, and was not full, so nothing to do @@ -311,26 +310,21 @@ impl Drop for Slab { fn drop(&mut self) { 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(); - + for chunk in self.chunks.iter() { #[cfg(test)] - std::eprintln!("Slab::drop({chunk_ptr:?}, layout: {layout:?})"); + std::eprintln!("Slab::drop({chunk:?}, layout: {layout:?})"); unsafe { - self.alloc.deallocate(chunk_ptr.cast(), layout); + self.alloc.deallocate(chunk.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(); + for chunk in self.full_chunks.iter() { + #[cfg(test)] + std::eprintln!("Slab::drop({chunk:?}, layout: {layout:?})"); unsafe { - self.alloc.deallocate(chunk_ptr.cast(), layout); + self.alloc.deallocate(chunk.cast(), layout); } } } @@ -438,7 +432,7 @@ impl SlabAllocator { mod tests { use super::*; - use core::{mem::forget, pin}; + use core::pin; use std::{prelude::rust_2024::*, sync::Mutex}; struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex>>);