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 liballoc::alloc::Allocator;
use crate::mem; use crate::{
collections::linked_list::{LinkedList, LinkedListNode},
mem,
};
const PAGE_SIZE: usize = 4096; const PAGE_SIZE: usize = 4096;
@ -14,9 +17,9 @@ pub struct Slab<A: Allocator + Clone> {
/// Size and alignment of each element in the slab. /// Size and alignment of each element in the slab.
element_size: usize, element_size: usize,
/// Pointer to the first chunk in the slab. /// 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. /// Pointer to the first chunk in the slab that is full.
full_head: Option<NonNull<SlabChunk<A>>>, full_chunks: LinkedList<SlabChunk<A>>,
alloc: A, alloc: A,
_pd: core::marker::PhantomPinned, _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 { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Slab") f.debug_struct("Slab")
.field("element_size", &self.element_size) .field("element_size", &self.element_size)
.field("head", &self.head) .finish_non_exhaustive()
.field("full_head", &self.full_head)
.finish()
} }
} }
#[derive(Debug)]
struct SlabChunk<A: Allocator + Clone> { struct SlabChunk<A: Allocator + Clone> {
/// Pointer to the next chunk in the slab. /// Pointer to the next chunk in the slab.
next: Option<NonNull<SlabChunk<A>>>, next: Option<NonNull<SlabChunk<A>>>,
@ -44,6 +46,16 @@ struct SlabChunk<A: Allocator + Clone> {
count: Cell<usize>, 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>>); struct ChunkSlot(Option<NonNull<Self>>);
enum SlotResult { enum SlotResult {
@ -109,8 +121,8 @@ impl<A: Allocator + Clone> Slab<A> {
let slab = Self { let slab = Self {
element_size, element_size,
head: None, chunks: LinkedList::new(),
full_head: None, full_chunks: LinkedList::new(),
alloc: alloc.clone(), alloc: alloc.clone(),
_pd: core::marker::PhantomPinned, _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 // 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. // 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!() panic!()
}; };
let chunk = bytes.as_non_null_ptr().cast::<SlabChunk<A>>(); let chunk = bytes.as_non_null_ptr().cast::<SlabChunk<A>>();
let _ = chunk.as_ptr().expose_provenance();
#[cfg(test)] #[cfg(test)]
std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})"); std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})");
@ -172,22 +185,22 @@ impl<A: Allocator + Clone> Slab<A> {
unsafe { unsafe {
let first_slot = chunk let first_slot = chunk
.as_ptr() .as_ptr()
.byte_add(Self::first_slot_offset(self.element_size)) .byte_add(Self::first_slot_offset(element_size))
.cast::<ChunkSlot>(); .cast::<ChunkSlot>();
for i in 0..(count - 1) { for i in 0..(count - 1) {
let chunk = first_slot.byte_add(i * self.element_size); let chunk = first_slot.byte_add(i * element_size);
let next = first_slot.byte_add((i + 1) * self.element_size); let next = first_slot.byte_add((i + 1) * element_size);
chunk.write(ChunkSlot(Some(NonNull::new_unchecked(next)))); chunk.write(ChunkSlot(Some(NonNull::new_unchecked(next))));
} }
first_slot first_slot
.byte_add((count - 1) * self.element_size) .byte_add((count - 1) * element_size)
.write(ChunkSlot(None)); .write(ChunkSlot(None));
chunk.write(SlabChunk { chunk.write(SlabChunk {
next: self.head, next: None,
// SAFETY: we only access `slab` via `SlabChunk::slab_pinned_mut` // 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))), free: Cell::new(Some(NonNull::new_unchecked(first_slot))),
count: Cell::new(0), count: Cell::new(0),
}); });
@ -196,36 +209,31 @@ impl<A: Allocator + Clone> Slab<A> {
chunk 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]> { fn alloc_slot(mut self: Pin<&mut Self>) -> NonNull<[u8]> {
#[cfg(test)] #[cfg(test)]
std::eprintln!("Slab::alloc_slot({self:?})"); std::eprintln!("Slab::alloc_slot({self:?})");
let chunk = match self.head { let slab: NonNull<Self> = unsafe { self.as_mut().get_unchecked_mut().into() };
Some(chunk) => chunk, let alloc = self.alloc.clone();
None => { let element_size = self.element_size;
let chunk = self.as_mut().alloc_chunk_cold();
unsafe { let mut chunk = unsafe {
self.as_mut().get_unchecked_mut().head = Some(chunk); self.as_mut()
} .get_unchecked_mut()
chunk .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::Some(non_null) => non_null,
SlotResult::Last(non_null) => { SlotResult::Last(non_null) => {
let chunk = chunk.remove();
unsafe { unsafe {
let mut_ref = self.as_mut().get_unchecked_mut(); self.as_mut()
mut_ref.head = chunk.next.take(); .map_unchecked_mut(|slab| &mut slab.full_chunks)
chunk.next = mut_ref.full_head; .push_front(chunk);
mut_ref.full_head = Some(chunk.into());
} }
non_null non_null
@ -241,11 +249,14 @@ impl<A: Allocator + Clone> Slab<A> {
fn free_slot(element_size: usize, slot: NonNull<u8>) { fn free_slot(element_size: usize, slot: NonNull<u8>) {
let (_, layout) = Self::count_and_layout(element_size); let (_, layout) = Self::count_and_layout(element_size);
let chunk_ptr = slot let addr =
.map_addr(|addr| unsafe { unsafe { NonZero::new_unchecked(mem::align_down(slot.addr().get(), layout.align())) };
NonZero::new_unchecked(mem::align_down(addr.get(), layout.align()))
}) // grab provenance exposed in alloc_chunk, since `slot` may be noalias restricted to its layout.
.cast::<SlabChunk<A>>();
// let chunk_ptr: NonNull<SlabChunk<A>> = NonNull::with_exposed_provenance(addr);
let chunk_ptr: NonNull<SlabChunk<A>> = slot.with_addr(addr).cast();
#[cfg(test)] #[cfg(test)]
std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})"); std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})");
@ -258,47 +269,35 @@ impl<A: Allocator + Clone> Slab<A> {
FreeSlotResult::Empty => { FreeSlotResult::Empty => {
// chunk is empty: // chunk is empty:
// first unlink the chunk from the slab.. // first unlink the chunk from the slab..
if slab.head == Some(chunk.into()) { if let Some(chunk) =
unsafe { slab.as_mut().get_unchecked_mut().head = chunk.next }; unsafe { slab.as_mut().map_unchecked_mut(|slab| &mut slab.chunks) }
} else { .remove_if(|c| c == chunk_ptr)
let mut head = slab.head.expect("chunk is linked, so head exists"); .next()
while let Some(next) = unsafe { (&raw const (*head.as_ptr()).next).read() } { {
if next == chunk.into() { // ..then free it.
unsafe { (&raw mut (*head.as_ptr()).next).write(chunk.next) }; #[cfg(test)]
break; std::eprintln!("free_slot::drop({chunk:?}, layout: {layout:?})");
}
head = next; 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 => { FreeSlotResult::WasFull => {
// chunk was full: // chunk was full:
// unlink it from the full list // unlink it from the full list
if slab.full_head == Some(chunk.into()) { if let Some(chunk) = unsafe {
unsafe { slab.as_mut().get_unchecked_mut().full_head = chunk.next }; slab.as_mut()
} else { .map_unchecked_mut(|slab| &mut slab.full_chunks)
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());
} }
.remove_if(|c| c == chunk_ptr)
.next()
{
unsafe {
slab.as_mut()
.map_unchecked_mut(|slab| &mut slab.chunks)
.push_front(chunk)
};
};
} }
FreeSlotResult::NotEmpty => { FreeSlotResult::NotEmpty => {
// chunk is not empty, and was not full, so nothing to do // 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) { fn drop(&mut self) {
let (_, layout) = Self::count_and_layout(self.element_size); let (_, layout) = Self::count_and_layout(self.element_size);
let mut chunk = self.head.take(); for chunk in self.chunks.iter() {
while let Some(mut chunk_ptr) = chunk {
let chunk_ref = unsafe { chunk_ptr.as_mut() };
chunk = chunk_ref.next.take();
#[cfg(test)] #[cfg(test)]
std::eprintln!("Slab::drop({chunk_ptr:?}, layout: {layout:?})"); std::eprintln!("Slab::drop({chunk:?}, layout: {layout:?})");
unsafe { unsafe {
self.alloc.deallocate(chunk_ptr.cast(), layout); self.alloc.deallocate(chunk.cast(), layout);
} }
} }
let mut chunk = self.full_head.take(); for chunk in self.full_chunks.iter() {
while let Some(mut chunk_ptr) = chunk { #[cfg(test)]
let chunk_ref = unsafe { chunk_ptr.as_mut() }; std::eprintln!("Slab::drop({chunk:?}, layout: {layout:?})");
chunk = chunk_ref.next.take();
unsafe { 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 { mod tests {
use super::*; use super::*;
use core::{mem::forget, pin}; use core::pin;
use std::{prelude::rust_2024::*, sync::Mutex}; use std::{prelude::rust_2024::*, sync::Mutex};
struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex<Pin<&'a mut SlabAllocator<A>>>); struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex<Pin<&'a mut SlabAllocator<A>>>);