janis/curiOS
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

slab alloc: use linked list

Browse source
This commit is contained in:
janis 2026-08-08 22:05:40 +02:00
parent 202cacd764
commit ef7f924143
Signed by: janis
SSH key fingerprint: SHA256:bB1qbbqmDXZNT0KKD5c2Dfjg53JGhj7B3CFcLIzSqq8
1 changed files with 83 additions and 89 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

172
crates/foundation/src/alloc/slab.rs
View file

@ -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<A: Allocator + Clone> {
/// Size and alignment of each element in the slab.
element_size: usize,
/// Pointer to the first chunk in the slab.
head: Option<NonNull<SlabChunk<A>>>,
chunks: LinkedList<SlabChunk<A>>,
/// Pointer to the first chunk in the slab that is full.
full_head: Option<NonNull<SlabChunk<A>>>,
full_chunks: LinkedList<SlabChunk<A>>,
alloc: A,
_pd: core::marker::PhantomPinned,
}
@ -25,12 +28,11 @@ 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()
.finish_non_exhaustive()
}
}
#[derive(Debug)]
struct SlabChunk<A: Allocator + Clone> {
/// Pointer to the next chunk in the slab.
next: Option<NonNull<SlabChunk<A>>>,
@ -44,6 +46,16 @@ struct SlabChunk<A: Allocator + Clone> {
count: Cell<usize>,
}
unsafe impl<A: Allocator + Clone> LinkedListNode for SlabChunk<A> {
fn next(this: NonNull<Self>) -> Option<NonNull<Self>> {
unsafe { (&raw const (*this.as_ptr()).next).read() }
}
fn set_next(this: NonNull<Self>, next: Option<NonNull<Self>>) {
unsafe { (&raw mut (*this.as_ptr()).next).write(next) };
}
}
struct ChunkSlot(Option<NonNull<Self>>);
enum SlotResult {
@ -109,8 +121,8 @@ impl<A: Allocator + Clone> Slab<A> {
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<A: Allocator + Clone> Slab<A> {
})
}
fn alloc_chunk(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
fn alloc_chunk(element_size: usize, alloc: &A, slab: NonNull<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(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::<SlabChunk<A>>();
let _ = chunk.as_ptr().expose_provenance();
#[cfg(test)]
std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})");
@ -172,22 +185,22 @@ impl<A: Allocator + Clone> Slab<A> {
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::<ChunkSlot>();
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<A: Allocator + Clone> Slab<A> {
chunk
}
#[cold]
fn alloc_chunk_cold(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
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<Self> = 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<A: Allocator + Clone> Slab<A> {
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 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<SlabChunk<A>> = NonNull::with_exposed_provenance(addr);
let chunk_ptr: NonNull<SlabChunk<A>> = slot.with_addr(addr).cast();
#[cfg(test)]
std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})");
@ -258,47 +269,35 @@ impl<A: Allocator + Clone> Slab<A> {
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<A: Allocator + Clone> Drop for Slab<A> {
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<A: Allocator + Clone> SlabAllocator<A> {
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<Pin<&'a mut SlabAllocator<A>>>);