kernel: pmm with overlapping check

num_traits dep

Cargo.lock

@@ -2,6 +2,12 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
 [[package]]
 name = "bit_field"
 version = "0.10.3"
@@ -14,13 +20,22 @@ version = "2.13.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da"
 
+[[package]]
+name = "foundation"
+version = "0.1.0"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "kernel"
 version = "0.1.0"
 dependencies = [
  "bit_field",
  "bitflags",
+ "foundation",
  "log",
+ "num-traits",
  "plain",
  "rbtree",
  "rustc-demangle",
@@ -33,6 +48,15 @@ version = "0.4.33"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
 
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "plain"
 version = "0.2.3"

crates/foundation/Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "foundation"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+num-traits = { version = "0.2.19", default-features = false }

crates/foundation/src/alloc/bump.rs

@@ -0,0 +1,426 @@
+
+//! A bump allocator inspired by / taken from the `stumpalo` crate
+
+use core::{
+    alloc::{Allocator, Layout},
+    cell::Cell,
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    ptr::{self, NonNull},
+};
+
+#[repr(align(16))]
+struct Chunk {
+    next: Option<NonNull<Chunk>>,
+    size: usize,
+}
+
+const EMPTY_CHUNK: Chunk = Chunk {
+    next: None,
+    size: 0,
+};
+
+struct RestorePoint {
+    top: *mut u8,
+    chunk: NonNull<Chunk>,
+}
+
+pub struct Bump<A: Allocator> {
+    top: Cell<*mut u8>,
+    bottom: Cell<*mut Chunk>,
+    next_chunk: Cell<Option<NonNull<Chunk>>>,
+    backing_alloc: A,
+}
+
+unsafe impl<A: Allocator + Send> Send for Bump<A> {}
+
+unsafe impl<A: Allocator> Allocator for BumpScope<'_, '_, A> {
+    fn allocate(
+        &self,
+        layout: Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
+        let virt = unsafe { self.alloc_layout(layout).ok_or(core::alloc::AllocError) }?;
+
+        Ok(virt.cast_slice(layout.size()))
+    }
+
+    unsafe fn deallocate(&self, _ptr: core::ptr::NonNull<u8>, _layout: Layout) {}
+}
+
+impl<A: Allocator> Bump<A> {
+    pub fn new_in(backing_alloc: A) -> Self {
+        let chunk = NonNull::from(&EMPTY_CHUNK);
+        let bottom = unsafe { chunk.as_ptr().add(1) };
+
+        Self {
+            top: Cell::new(bottom.cast::<u8>()),
+            bottom: Cell::new(bottom),
+            next_chunk: Cell::new(None),
+            backing_alloc,
+        }
+    }
+
+    /// # Safety
+    /// The caller must ensure that the chunk is aligned to `Chunk` alignment.
+    pub unsafe fn from_raw_chunk_in(chunk: NonNull<[u8]>, backing_alloc: A) -> Self {
+        assert!(
+            chunk.as_ptr().addr().is_multiple_of(align_of::<Chunk>()),
+            "Chunk must be aligned to Chunk alignment"
+        );
+        assert!(
+            chunk.len() >= size_of::<Chunk>(),
+            "Chunk must be at least the size of Chunk"
+        );
+
+        let len = chunk.len();
+        let chunk = chunk.cast::<Chunk>();
+        let bottom = unsafe { chunk.as_ptr().add(1) };
+        let top = unsafe { bottom.byte_add(len - size_of::<Chunk>()) };
+
+        Self {
+            top: Cell::new(top.cast::<u8>()),
+            bottom: Cell::new(bottom),
+            next_chunk: Cell::new(None),
+            backing_alloc,
+        }
+    }
+
+    pub fn as_scope<'env>(&'env mut self) -> BumpScope<'env, 'env, A> {
+        BumpScope {
+            bump: self,
+            _env: PhantomData,
+            _scope: PhantomData,
+        }
+    }
+
+    pub fn scope<'env, F, R>(&'env mut self, f: F) -> R
+    where
+        F: for<'scope> FnOnce(&'scope mut BumpScope<'env, 'scope, A>) -> R + 'env,
+        R: 'env,
+    {
+        let restore = self.restore_point();
+
+        let mut scope = BumpScope {
+            bump: self,
+            _env: PhantomData,
+            _scope: PhantomData,
+        };
+
+        let result = f(&mut scope);
+
+        unsafe { self.restore(restore) };
+
+        result
+    }
+
+    fn restore_point(&self) -> RestorePoint {
+        RestorePoint {
+            top: self.top.get(),
+            chunk: NonNull::from(self.chunk()),
+        }
+    }
+
+    unsafe fn restore(&self, restore_point: RestorePoint) {
+        // get current chunk and next chunk
+        let mut chunk = Some(NonNull::from(self.chunk()));
+        let mut head = self.next_chunk.get();
+
+        // walk the linked list of chunks used since the restore point and
+        // re-link them onto the free list.
+        while let Some(cnk) = chunk
+            && cnk != restore_point.chunk
+        {
+            let next = unsafe { ptr::replace(&raw mut (*cnk.as_ptr()).next, head) };
+
+            head = Some(cnk);
+            chunk = next;
+        }
+
+        self.top.set(restore_point.top);
+        let bot = unsafe { restore_point.chunk.as_ptr().add(1) };
+        self.bottom.set(bot);
+        self.next_chunk.set(head);
+    }
+
+    #[inline]
+    pub unsafe fn alloc_layout(&self, layout: Layout) -> Option<NonNull<u8>> {
+        let top = self.top.get();
+        let bottom = self.bottom.get().addr();
+        let extra = Self::extra_bytes(top, layout.align());
+
+        let slow_path = Self::out_of_mem(layout, top, bottom, false);
+
+        if slow_path {
+            self.alloc_layout_slow_cold(layout)
+        } else {
+            let new_top = unsafe { top.byte_sub(extra).byte_sub(layout.size()) };
+            self.top.set(new_top);
+
+            Some(unsafe { NonNull::new_unchecked(new_top) })
+        }
+    }
+
+    #[inline(never)]
+    fn try_alloc_slow_with_no_inline<F, T>(&self, f: F) -> Option<NonNull<T>>
+    where
+        F: FnOnce() -> T,
+    {
+        let p = self.alloc_layout_slow(Layout::new::<T>())?;
+        let p = p.cast::<T>();
+        unsafe { p.write(f()) };
+        Some(p)
+    }
+
+    #[cold]
+    fn alloc_layout_slow_cold(&self, layout: Layout) -> Option<NonNull<u8>> {
+        self.alloc_layout_slow(layout)
+    }
+
+    // #[inline(never)]
+    // fn alloc_layout_slow_no_inline(&self, layout: Layout) -> Option<NonNull<u8>> {
+    //     self.alloc_layout_slow(layout)
+    // }
+
+    fn alloc_layout_slow(&self, layout: Layout) -> Option<NonNull<u8>> {
+        let extra_with_chunk = {
+            assert!(size_of::<Chunk>() == align_of::<Chunk>());
+            // if the allocation is greater than the size of the chunk header, we need more bytes past the chunk header to align the allocation.
+            if layout.align() > size_of::<Chunk>() {
+                layout.align() - size_of::<Chunk>()
+            } else {
+                0
+            }
+        };
+
+        let min_size = layout.size() + extra_with_chunk;
+
+        while let Some(mut chunk) = self.next_chunk.get() {
+            // SAFETY: `Bump` is not `Sync` and `alloc_layout_slow is not
+            // reentrant, we have exclusive access to `chunk`
+            let chunk = unsafe { chunk.as_mut() };
+            let cap = chunk.size;
+
+            if cap >= min_size {
+                let prev = self.chunk();
+                self.next_chunk.set(chunk.next);
+
+                // add the current chunk to the linked list of chunks, so it can be freed later.
+                chunk.next = Some(NonNull::from(prev));
+
+                unsafe {
+                    // bottom is just past the chunk header
+                    let bot = (&raw const *chunk).add(1).cast_mut();
+                    self.bottom.set(bot);
+
+                    // calculate top of chunk
+                    let top = bot.byte_add(cap).cast::<u8>();
+
+                    // sub allocation
+                    let extra = Self::extra_bytes(top, layout.align());
+                    let top = top.byte_sub(extra).byte_sub(layout.size());
+                    self.top.set(top);
+
+                    return Some(NonNull::new_unchecked(top));
+                };
+            } else {
+                // chunk is too small, try the next one
+                self.next_chunk.set(chunk.next);
+                unsafe {
+                    let layout = Layout::from_size_align_unchecked(
+                        cap + size_of::<Chunk>(),
+                        align_of::<Chunk>(),
+                    );
+                    self.backing_alloc
+                        .deallocate(NonNull::from(chunk).cast(), layout);
+                }
+            }
+        }
+
+        let chunk = self.chunk();
+        let prev = if chunk.size == 0 {
+            None
+        } else {
+            Some(NonNull::from(chunk))
+        };
+
+        let size = self.next_chunk_size().max(min_size);
+        let new_chunk = self
+            .backing_alloc
+            .allocate(unsafe {
+                Layout::from_size_align_unchecked(size + size_of::<Chunk>(), align_of::<Chunk>())
+            })
+            .ok()?
+            .cast::<Chunk>();
+
+        unsafe {
+            new_chunk.write(Chunk { next: prev, size });
+
+            let bot = new_chunk.add(1).as_ptr();
+            self.bottom.set(bot);
+
+            let top = bot.byte_add(size).cast::<u8>();
+            let extra = Self::extra_bytes(top, layout.align());
+            let top = top.byte_sub(extra).byte_sub(layout.size());
+            self.top.set(top);
+
+            Some(NonNull::new_unchecked(top))
+        }
+    }
+
+    fn next_chunk_size(&self) -> usize {
+        let chunk = self.chunk();
+        let cap = (chunk.size + size_of::<Chunk>()).saturating_mul(2) - size_of::<Chunk>();
+
+        cap.max(0x1000)
+    }
+
+    pub unsafe fn alloc_raw<T>(&self) -> Option<NonNull<T>> {
+        let layout = core::alloc::Layout::new::<T>();
+        unsafe { self.alloc_layout(layout).map(|p| p.cast::<T>()) }
+    }
+
+    fn out_of_mem(layout: Layout, top: *mut u8, bottom: usize, comptime: bool) -> bool {
+        let extra = Self::extra_bytes(top, layout.align());
+
+        let max_padding = layout.align() - 1;
+        let max_size = layout.size() + max_padding;
+        let top = top.addr();
+
+        // biggest possible virtual address on x86_64 is 2^57 - 1
+        const MAX_ADDR: usize = ((1u64 << 57) - 1) as usize;
+        // offsets of < SAFE_SIZE are guaranteed not to overflow the address space.
+        const SAFE_SIZE: usize = usize::MAX - MAX_ADDR;
+
+        if comptime && max_size < SAFE_SIZE {
+            if max_size <= 16 {
+                let top = top - extra;
+                let top = top - layout.size();
+                bottom > top
+            } else if max_padding < 16 {
+                let top = top - extra;
+                bottom + layout.size() > top
+            } else {
+                bottom + layout.size() + max_padding > top
+            }
+        } else if max_padding
+            .checked_add(isize::MAX as usize)
+            .is_some_and(|p| p < SAFE_SIZE)
+        {
+            if max_size < 16 {
+                let top = top - extra;
+                bottom + layout.size() > top
+            } else {
+                bottom + extra + layout.size() > top
+            }
+        } else {
+            if max_padding < 16 {
+                let top = top - extra;
+                top.checked_sub(layout.size()).is_none_or(|t| t < bottom)
+            } else {
+                top.checked_sub(extra)
+                    .and_then(|t| t.checked_sub(layout.size()))
+                    .is_none_or(|t| t < bottom)
+            }
+        }
+    }
+
+    fn extra_bytes(ptr: *mut u8, align: usize) -> usize {
+        assert!(align.is_power_of_two(), "Alignment must be a power of two");
+        ptr.addr() & (align - 1)
+    }
+
+    fn chunk(&self) -> &Chunk {
+        unsafe { self.bottom.get().sub(1).as_ref_unchecked() }
+    }
+}
+
+#[repr(transparent)]
+pub struct BumpScope<'env, 'scope, A: Allocator> {
+    bump: &'scope mut Bump<A>,
+    _env: PhantomData<&'env &'env mut ()>,
+    _scope: PhantomData<&'scope &'scope mut ()>,
+}
+
+impl<'env, 'scope, A: Allocator> Deref for BumpScope<'env, 'scope, A> {
+    type Target = Bump<A>;
+
+    fn deref(&self) -> &Self::Target {
+        self.bump
+    }
+}
+
+impl<'env, 'scope, A: Allocator> DerefMut for BumpScope<'env, 'scope, A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.bump
+    }
+}
+
+impl<'env, 'scope, A: Allocator> BumpScope<'env, 'scope, A> {
+    pub fn alloc_with<F, T>(&self, f: F) -> &'scope mut T
+    where
+        F: FnOnce() -> T,
+    {
+        self.try_alloc_with(f).unwrap_or_else(|| {
+            panic!(
+                "Bump allocator out of memory when allocating {} bytes with alignment {}",
+                size_of::<T>(),
+                align_of::<T>()
+            )
+        })
+    }
+
+    pub fn try_alloc_with<F, T>(&self, f: F) -> Option<&'scope mut T>
+    where
+        F: FnOnce() -> T,
+    {
+        let layout = Layout::new::<T>();
+        let top = self.top.get();
+        let bottom = self.bottom.get().addr();
+        let extra = Bump::<A>::extra_bytes(top, layout.align());
+
+        let slow_path = Bump::<A>::out_of_mem(layout, top, bottom, false);
+
+        let mut ptr = if layout.size() <= 16 {
+            if slow_path {
+                self.try_alloc_slow_with_no_inline(f)?
+            } else {
+                let new_top = unsafe { top.byte_sub(extra).byte_sub(layout.size()) };
+                self.top.set(new_top);
+                unsafe {
+                    ptr::write(new_top.cast::<T>(), f());
+                    NonNull::new_unchecked(new_top.cast::<T>())
+                }
+            }
+        } else {
+            // in-place
+            let ptr = if slow_path {
+                self.alloc_layout_slow_cold(layout)?
+            } else {
+                let new_top = unsafe { top.byte_sub(extra).byte_sub(layout.size()) };
+                self.top.set(new_top);
+
+                unsafe { NonNull::new_unchecked(new_top) }
+            };
+
+            let ptr = ptr.cast::<T>();
+            unsafe {
+                ptr.write(f());
+            }
+
+            ptr
+        };
+
+        Some(unsafe { ptr.as_mut() })
+    }
+}
+
+fn _asdf(bump: &mut Bump<super::PanicingAllocator>) {
+    let mut bump = bump.as_scope();
+    let x = bump.alloc_with(|| 3u64);
+    bump.scope(|bump| {
+        let _y = bump.alloc_with(|| 4u64);
+        let _z = bump.alloc_with(|| 5u64);
+    });
+
+    assert!(*x == 3);
+}

crates/foundation/src/alloc/mod.rs

@@ -0,0 +1,23 @@
+use core::ptr::NonNull;
+
+use liballoc::alloc::Allocator;
+
+pub mod bump;
+pub mod slab;
+
+#[derive(Debug, Default)]
+pub struct PanicingAllocator;
+
+unsafe impl Allocator for PanicingAllocator {
+    fn allocate(
+        &self,
+        layout: core::alloc::Layout,
+    ) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
+        match layout.size() {
+            0 => Ok(NonNull::slice_from_raw_parts(NonNull::dangling(), 0)),
+            _ => panic!("PanicingAllocator cannot allocate memory"),
+        }
+    }
+
+    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: core::alloc::Layout) {}
+}

crates/foundation/src/alloc/slab.rs

@@ -0,0 +1,483 @@
+//! A slab allocator
+
+use core::{
+    alloc::Layout, cell::Cell, fmt::Debug, hint::unlikely, num::NonZero, pin::Pin, ptr::NonNull,
+};
+
+use liballoc::alloc::Allocator;
+
+use crate::mem;
+
+const PAGE_SIZE: usize = 4096;
+
+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>>>,
+    /// 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> {
+    /// Pointer to the next chunk in the slab.
+    next: Option<NonNull<SlabChunk<A>>>,
+    /// Pointer to the slab that owns this chunk.
+    slab: NonNull<Slab<A>>,
+    /// Linked list of free elements in this slab. When this is `None`, the
+    /// slab is full.
+    free: Cell<Option<NonNull<ChunkSlot>>>,
+    /// Number of outstanding allocations from this slab. When this reaches
+    /// zero, the slab can be freed.
+    count: Cell<usize>,
+}
+
+struct ChunkSlot(Option<NonNull<Self>>);
+
+enum SlotResult {
+    Some(NonNull<u8>),
+    Last(NonNull<u8>),
+    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;
+        };
+
+        self.free.set(unsafe { slot.as_ref() }.next());
+        self.count.update(|count| count + 1);
+
+        match self.free.get() {
+            Some(_) => SlotResult::Some(slot.cast()),
+            None => SlotResult::Last(slot.cast()),
+        }
+    }
+
+    /// returns `true` if the slab is now empty and can be freed
+    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);
+
+        match (next, self.count.get()) {
+            (None, _) => FreeSlotResult::WasFull,
+            (_, 0) => FreeSlotResult::Empty,
+            _ => FreeSlotResult::NotEmpty,
+        }
+    }
+}
+
+impl ChunkSlot {
+    fn next(&self) -> Option<NonNull<Self>> {
+        self.0
+    }
+}
+
+impl<A: Allocator + Clone> Slab<A> {
+    fn new(element_size: usize, alloc: A) -> Self {
+        assert!(
+            element_size.is_power_of_two(),
+            "element_size must be a power of two"
+        );
+
+        let slab = Self {
+            element_size,
+            head: None,
+            full_head: None,
+            alloc: alloc.clone(),
+            _pd: core::marker::PhantomPinned,
+        };
+
+        #[cfg(test)]
+        std::eprintln!("Slab::new({element_size}) -> {slab:#?}",);
+
+        slab
+    }
+
+    fn first_slot_offset(element_size: usize) -> usize {
+        mem::align_up(core::mem::size_of::<SlabChunk<A>>(), element_size)
+    }
+
+    fn count_and_layout(element_size: usize) -> (usize, Layout) {
+        let (count, size, align) = {
+            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(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));
+
+                (count, size, size)
+            } else {
+                let count = one_page_count;
+                let size = PAGE_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(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(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(self.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);
+                chunk.write(ChunkSlot(Some(NonNull::new_unchecked(next))));
+            }
+            first_slot
+                .byte_add((count - 1) * self.element_size)
+                .write(ChunkSlot(None));
+
+            chunk.write(SlabChunk {
+                next: self.head,
+                // 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),
+            });
+        }
+
+        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 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) => {
+                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 => {
+                panic!("SlabChunk is full, but it is still the head of the slab");
+            }
+        };
+
+        ptr.cast_slice(self.element_size)
+    }
+
+    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>>();
+
+        #[cfg(test)]
+        std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})");
+
+        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 { (&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 {
+                    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(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);
+            }
+        }
+    }
+}
+
+const SLAB_ALLOCATOR_BUCKETS: usize = 8;
+pub struct SlabAllocator<A: Allocator + Clone> {
+    /// Slabs for each power-of-two from 16 bytes to 2048 bytes (inclusive).
+    slabs: [Slab<A>; SLAB_ALLOCATOR_BUCKETS],
+    alloc: A,
+}
+
+impl<A: Allocator + Clone> SlabAllocator<A> {
+    pub fn new(alloc: A) -> Self {
+        Self {
+            slabs: [
+                Slab::new(16, alloc.clone()),
+                Slab::new(32, alloc.clone()),
+                Slab::new(64, alloc.clone()),
+                Slab::new(128, alloc.clone()),
+                Slab::new(256, alloc.clone()),
+                Slab::new(512, alloc.clone()),
+                Slab::new(1024, alloc.clone()),
+                Slab::new(2048, alloc.clone()),
+            ],
+            alloc,
+        }
+    }
+
+    fn slab_index_for_size(size: usize) -> Option<usize> {
+        if size == 0 {
+            return None;
+        }
+
+        // the smallest slab is 16 bytes
+        let size = size.max(16);
+
+        // get the index of the slab by calculating the log2 of the size and
+        // 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)
+        } else {
+            None
+        }
+    }
+
+    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));
+        }
+
+        let size = layout.size().max(layout.align());
+
+        match Self::slab_index_for_size(size) {
+            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()
+            }
+        }
+    }
+
+    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;
+        }
+
+        let size = layout.size().max(layout.align());
+
+        match Self::slab_index_for_size(size) {
+            Some(slab_index) => unsafe {
+                Slab::<A>::free_slot(self.slabs.get_unchecked(slab_index).element_size, ptr);
+            },
+            None => {
+                // deallocate directly to the backing allocator
+                unsafe { self.alloc.deallocate(ptr, layout) }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use core::{mem::forget, pin};
+    use std::{prelude::rust_2024::*, sync::Mutex};
+
+    struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex<Pin<&'a mut SlabAllocator<A>>>);
+
+    unsafe impl<'a, A: Allocator + Clone> Allocator for SlabAllocatorWrapper<'a, A> {
+        fn allocate(
+            &self,
+            layout: core::alloc::Layout,
+        ) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
+            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().as_mut().dealloc(ptr, layout)
+        }
+    }
+
+    #[test]
+    fn test_slab_allocator() {
+        let mut slab = pin::pin!(SlabAllocator::new(std::alloc::Global));
+        let alloc = SlabAllocatorWrapper(Mutex::new(slab.as_mut()));
+
+        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) };
+            drop(boxed);
+            std::eprintln!("test_slab_allocator: dropped boxed");
+        }
+    }
+}

crates/foundation/src/lib.rs

@@ -0,0 +1,99 @@
+#![no_std]
+#![feature(
+    const_trait_impl,
+    const_default,
+    allocator_api,
+    ptr_cast_slice,
+    likely_unlikely,
+    slice_ptr_get
+)]
+#![cfg_attr(test, feature(box_as_non_null))]
+
+extern crate alloc as liballoc;
+
+#[cfg(test)]
+extern crate std;
+
+pub mod alloc;
+pub mod sync;
+pub mod mem {
+    use num_traits::PrimInt;
+
+    pub fn replace<T, R>(v: &mut T, f: impl FnOnce(T) -> (T, R)) -> R {
+        struct Guard;
+        impl Drop for Guard {
+            fn drop(&mut self) {
+                panic!("replace() panicked");
+            }
+        }
+        let guard = Guard;
+        let value = unsafe { core::ptr::read(v) };
+        let (new_value, ret) = f(value);
+        unsafe { core::ptr::write(v, new_value) };
+        core::mem::forget(guard);
+        ret
+    }
+
+    pub fn align_up<T: PrimInt>(value: T, alignment: T) -> T {
+        (value + alignment - T::one()) & !(alignment - T::one())
+    }
+
+    pub fn align_down<T: PrimInt>(value: T, alignment: T) -> T {
+        value & !(alignment - T::one())
+    }
+
+    /// # Safety
+    pub unsafe fn volatile_copy<T: Sized>(src: *const T, dst: *mut T, count: usize) {
+        unsafe {
+            for i in 0..count {
+                let src_ptr = src.add(i);
+                let dst_ptr = dst.add(i);
+                let val = src_ptr.read();
+                dst_ptr.write_volatile(val);
+            }
+        }
+    }
+}
+
+pub struct DropGuard<F: FnOnce()>(::core::mem::ManuallyDrop<F>);
+impl<F: FnOnce()> DropGuard<F> {
+    pub fn new(f: F) -> Self {
+        DropGuard(::core::mem::ManuallyDrop::new(f))
+    }
+    #[allow(dead_code)]
+    pub fn forget(self) {
+        let mut this = ::core::mem::ManuallyDrop::new(self);
+        unsafe {
+            ::core::mem::ManuallyDrop::drop(&mut this.0);
+        }
+    }
+}
+
+impl<F: FnOnce()> Drop for DropGuard<F> {
+    fn drop(&mut self) {
+        unsafe { ::core::ptr::read(&*self.0)() }
+    }
+}
+
+#[derive(Default)]
+pub struct PanicGuard;
+
+impl PanicGuard {
+    pub fn new() -> Self {
+        Self
+    }
+
+    pub fn forget(self) {
+        core::mem::forget(self);
+    }
+}
+
+impl Drop for PanicGuard {
+    #[track_caller]
+    fn drop(&mut self) {
+        panic!(
+            "PanicGuard dropped without being forgotten: {}",
+            core::panic::Location::caller()
+        );
+    }
+}

crates/foundation/src/sync/mod.rs

@@ -0,0 +1,487 @@
+mod spin_mutex {
+    use core::{
+        cell::UnsafeCell,
+        sync::atomic::{AtomicBool, Ordering},
+    };
+
+    pub struct SpinMutex<T: ?Sized> {
+        lock: AtomicBool,
+        data: UnsafeCell<T>,
+    }
+
+    pub struct SpinMutexGuard<'a, T: ?Sized + 'a> {
+        mutex: &'a SpinMutex<T>,
+    }
+
+    impl<T: ?Sized> core::ops::Deref for SpinMutexGuard<'_, T> {
+        type Target = T;
+
+        fn deref(&self) -> &Self::Target {
+            unsafe { &*self.mutex.data.get() }
+        }
+    }
+
+    impl<T: ?Sized> core::ops::DerefMut for SpinMutexGuard<'_, T> {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            unsafe { &mut *self.mutex.data.get() }
+        }
+    }
+
+    impl<T: ?Sized> Drop for SpinMutexGuard<'_, T> {
+        fn drop(&mut self) {
+            unsafe { self.mutex.unlock() }
+        }
+    }
+
+    unsafe impl<T: ?Sized + Send> Send for SpinMutex<T> {}
+    // SAFETY: SpinMutex properly synchronises access to T, but T must be Send
+    // to be safely shared between threads.
+    unsafe impl<T: ?Sized + Send> Sync for SpinMutex<T> {}
+
+    // SAFETY: &SpinMutexGuard cannot yield a &mut T, and the SpinMutexGuard
+    // holds the lock for the duration of its lifetime, so it is safe to share
+    // between threads.
+    unsafe impl<T: ?Sized + Sync> Sync for SpinMutexGuard<'_, T> {}
+    // SAFETY: SpinMutex may be unlocked on any thread, so it is safe to send
+    // the guard to another thread as long as T is Send.
+    unsafe impl<T: ?Sized + Send> Send for SpinMutexGuard<'_, T> {}
+
+    impl<T> SpinMutex<T> {
+        pub const fn new(data: T) -> Self {
+            Self {
+                lock: AtomicBool::new(false),
+                data: UnsafeCell::new(data),
+            }
+        }
+
+        pub fn into_inner(self) -> T {
+            self.data.into_inner()
+        }
+    }
+
+    impl<T: ?Sized> SpinMutex<T> {
+        pub fn as_mut_ptr(&self) -> *mut T {
+            self.data.get()
+        }
+
+        /// # Safety
+        /// The caller must be the logical owner of the locked mutex.
+        #[expect(clippy::mut_from_ref, reason = "unsafe mutex api")]
+        pub unsafe fn get_mut_unchecked(&self) -> &mut T {
+            unsafe { &mut *self.data.get() }
+        }
+
+        /// # Safety
+        /// The caller must be the logical owner of the locked mutex.
+        pub unsafe fn get_unchecked(&self) -> &T {
+            unsafe { &*self.data.get() }
+        }
+
+        pub fn try_lock(&self) -> bool {
+            self.lock
+                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+                .is_ok()
+        }
+        pub fn try_lock_weak(&self) -> bool {
+            self.lock
+                .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
+                .is_ok()
+        }
+        pub fn lock(&self) -> SpinMutexGuard<'_, T> {
+            while !self.try_lock() {
+                core::hint::spin_loop();
+            }
+            SpinMutexGuard { mutex: self }
+        }
+
+        /// # Safety
+        /// The caller must be the logical owner of the locked mutex.
+        pub unsafe fn unlock(&self) {
+            self.lock.store(false, Ordering::Release);
+        }
+    }
+}
+
+pub use spin_mutex::SpinMutex;
+
+mod once {
+    use core::{
+        cell::{Cell, UnsafeCell},
+        mem::{ManuallyDrop, MaybeUninit},
+        sync::atomic::{AtomicU8, Ordering},
+    };
+
+    const UNINITIALIZED: u8 = 0;
+    const INITIALIZING: u8 = 1;
+    const INITIALIZED: u8 = 2;
+    const POISONED: u8 = 3;
+
+    pub struct OnceState {
+        poisoned: bool,
+        state_to_set: Cell<u8>,
+    }
+
+    impl OnceState {
+        pub fn is_poisoned(&self) -> bool {
+            self.poisoned
+        }
+        pub fn poison(&self) {
+            self.state_to_set.set(POISONED);
+        }
+    }
+
+    pub struct Once {
+        state: AtomicU8,
+    }
+
+    const impl Default for Once {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+
+    impl Once {
+        pub const fn new() -> Self {
+            Self {
+                state: AtomicU8::new(UNINITIALIZED),
+            }
+        }
+
+        pub fn state(&self) -> u8 {
+            self.state.load(Ordering::Acquire)
+        }
+
+        pub fn set_state(&self, state: u8) {
+            self.state.store(state, Ordering::Release);
+        }
+
+        pub fn is_completed(&self) -> bool {
+            self.state.load(Ordering::Acquire) == INITIALIZED
+        }
+
+        pub fn call_once<F>(&self, f: F)
+        where
+            F: FnOnce(&OnceState),
+        {
+            if self.is_completed() {
+                return;
+            }
+
+            self.call_once_slow(false, f);
+        }
+
+        pub fn call_once_force<F>(&self, f: F)
+        where
+            F: FnOnce(&OnceState),
+        {
+            if self.is_completed() {
+                return;
+            }
+
+            self.call_once_slow(true, f);
+        }
+
+        #[cold]
+        pub fn call_once_slow<F>(&self, ignore_poison: bool, f: F)
+        where
+            F: FnOnce(&OnceState),
+        {
+            let mut state = self.state.load(Ordering::Acquire);
+            loop {
+                match state {
+                    INITIALIZED => return,
+                    POISONED if !ignore_poison => {
+                        panic!("Once instance has previously been poisoned")
+                    }
+                    POISONED | UNINITIALIZED => {
+                        match self.state.compare_exchange(
+                            state,
+                            INITIALIZING,
+                            Ordering::Acquire,
+                            // if we get `Err(INITIALIZED)`, we want to have
+                            // acquired what the lock is protecting.
+                            Ordering::Acquire,
+                        ) {
+                            Err(new) => {
+                                state = new;
+                                continue;
+                            }
+                            Ok(_) => {
+                                // even though we don't have unwinding, for
+                                // completeness sake we'll poison the lock if
+                                // the closure panics.
+                                let guard = crate::DropGuard::new(|| {
+                                    self.state.store(POISONED, Ordering::Release)
+                                });
+
+                                let state = OnceState {
+                                    poisoned: state == POISONED,
+                                    state_to_set: Cell::new(INITIALIZED),
+                                };
+
+                                f(&state);
+                                guard.forget();
+
+                                self.state
+                                    .store(state.state_to_set.take(), Ordering::Release);
+                                return;
+                            }
+                        }
+                    }
+                    _ => {
+                        assert_eq!(state, INITIALIZING);
+                        loop {
+                            state = self.state.load(Ordering::Acquire);
+                            if state != INITIALIZING {
+                                break;
+                            }
+                            core::hint::spin_loop();
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    pub struct OnceLock<T> {
+        once: Once,
+        data: UnsafeCell<MaybeUninit<T>>,
+    }
+
+    unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
+
+    const impl<T> Default for OnceLock<T> {
+        fn default() -> Self {
+            Self {
+                once: Once::new(),
+                data: UnsafeCell::new(MaybeUninit::uninit()),
+            }
+        }
+    }
+
+    impl<T> OnceLock<T> {
+        pub const fn new() -> Self {
+            Self {
+                once: Once::new(),
+                data: UnsafeCell::new(MaybeUninit::uninit()),
+            }
+        }
+
+        pub const fn from(t: T) -> Self {
+            Self {
+                once: Once::new(),
+                data: UnsafeCell::new(MaybeUninit::new(t)),
+            }
+        }
+
+        /// # Safety
+        /// The caller must ensure that the `OnceLock` has been initialized before calling this method.
+        pub unsafe fn get_unchecked(&self) -> &T {
+            unsafe { self.data.get().as_ref_unchecked().assume_init_ref() }
+        }
+
+        /// # Safety
+        /// The caller must ensure that the `OnceLock` has been initialized before calling this method.
+        pub unsafe fn get_mut_unchecked(&mut self) -> &mut T {
+            unsafe { self.data.get().as_mut_unchecked().assume_init_mut() }
+        }
+
+        pub fn get(&self) -> Option<&T> {
+            if self.is_completed() {
+                Some(unsafe { self.get_unchecked() })
+            } else {
+                None
+            }
+        }
+
+        pub fn get_mut(&mut self) -> Option<&mut T> {
+            if self.is_completed_mut() {
+                Some(unsafe { self.get_mut_unchecked() })
+            } else {
+                None
+            }
+        }
+
+        pub fn is_completed(&self) -> bool {
+            self.once.is_completed()
+        }
+
+        pub fn is_completed_mut(&mut self) -> bool {
+            // avoid atomic load when we have mutable access to self
+            *self.once.state.get_mut() == INITIALIZED
+        }
+
+        pub fn get_spinning(&self) -> Option<&T> {
+            loop {
+                match self.once.state() {
+                    INITIALIZED => return Some(unsafe { self.get_unchecked() }),
+                    INITIALIZING => core::hint::spin_loop(),
+                    UNINITIALIZED | POISONED => return None,
+                    _ => unreachable!(),
+                }
+            }
+        }
+
+        pub fn initialize<F, E>(&self, f: F) -> Result<(), E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            let mut res: Result<(), E> = Ok(());
+            let cell = unsafe { &mut *self.data.get() };
+
+            self.once.call_once_force(|state| match f() {
+                Ok(value) => {
+                    unsafe { cell.as_mut_ptr().write(value) };
+                }
+                Err(e) => {
+                    state.poison();
+                    res = Err(e);
+                }
+            });
+
+            res
+        }
+
+        pub fn get_or_init<F>(&self, f: F) -> &T
+        where
+            F: FnOnce() -> T,
+        {
+            match self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
+                Ok(value) => value,
+                Err(_) => unreachable!(),
+            }
+        }
+
+        pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
+        where
+            F: FnOnce() -> T,
+        {
+            match self.get_mut_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
+                Ok(value) => value,
+                Err(_) => unreachable!(),
+            }
+        }
+
+        pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            if !self.is_completed() {
+                self.initialize(f)?;
+            }
+
+            unsafe { Ok(self.get_unchecked()) }
+        }
+
+        pub fn get_mut_or_try_init<F, E>(&mut self, f: F) -> Result<&mut T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            if self.get_mut().is_none() {
+                self.initialize(f)?;
+            }
+
+            Ok(unsafe { self.get_mut_unchecked() })
+        }
+
+        pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
+            let mut val = Some(value);
+            let res = self.get_or_init(|| unsafe { val.take().unwrap_unchecked() });
+            match val {
+                Some(value) => Err((res, value)),
+                None => Ok(res),
+            }
+        }
+    }
+
+    impl<T> Drop for OnceLock<T> {
+        fn drop(&mut self) {
+            if self.is_completed() {
+                unsafe { self.data.get().as_mut_unchecked().assume_init_drop() }
+            }
+        }
+    }
+
+    pub struct LazyLock<T, F = fn() -> T> {
+        once: Once,
+        f: UnsafeCell<ManuallyDrop<F>>,
+        t: UnsafeCell<MaybeUninit<T>>,
+    }
+
+    unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
+
+    impl<T, F> LazyLock<T, F>
+    where
+        F: FnOnce() -> T,
+    {
+        pub const fn new(f: F) -> Self {
+            Self {
+                once: Once::new(),
+                f: UnsafeCell::new(ManuallyDrop::new(f)),
+                t: UnsafeCell::new(MaybeUninit::uninit()),
+            }
+        }
+
+        pub fn get(&self) -> Option<&T> {
+            if self.once.is_completed() {
+                Some(unsafe { (&*self.t.get()).assume_init_ref() })
+            } else {
+                None
+            }
+        }
+
+        fn force(this: &Self) -> &T {
+            this.once.call_once(|_| {
+                // SAFETY: because `call_once` will panic if poisoned, this
+                // closure will only be called once.
+                let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) };
+                let val = f();
+                unsafe { (&mut *this.t.get()).write(val) };
+            });
+
+            unsafe { (&*this.t.get()).assume_init_ref() }
+        }
+
+        fn force_mut(this: &mut Self) -> &mut T {
+            this.once.call_once(|_| {
+                // SAFETY: because `call_once` will panic if poisoned, this
+                // closure will only be called once.
+                let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) };
+                let val = f();
+                unsafe { (&mut *this.t.get()).write(val) };
+            });
+
+            unsafe { (&mut *this.t.get()).assume_init_mut() }
+        }
+    }
+
+    impl<T, F: FnOnce() -> T> core::ops::Deref for LazyLock<T, F> {
+        type Target = T;
+
+        fn deref(&self) -> &Self::Target {
+            LazyLock::force(self)
+        }
+    }
+
+    impl<T, F: FnOnce() -> T> core::ops::DerefMut for LazyLock<T, F> {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            LazyLock::force_mut(self)
+        }
+    }
+
+    impl<T, F> Drop for LazyLock<T, F> {
+        fn drop(&mut self) {
+            match self.once.state() {
+                UNINITIALIZED => unsafe {
+                    ManuallyDrop::drop(self.f.get_mut());
+                },
+                INITIALIZED => unsafe {
+                    (&mut *self.t.get()).assume_init_drop();
+                },
+                _ => {}
+            }
+        }
+    }
+}
+
+pub use once::{LazyLock, Once, OnceLock};

crates/rbtree/src/lib.rs

@@ -6,4 +6,22 @@ mod raw_node;
 
 extern crate alloc;
 
-pub use raw_node::{Color, RBTree, Side, TreeIter, TreeNodeIter, UnsafeNode};
+pub use raw_node::{
+    Color, Drain, Handle, LeftOrRight, RBTree, SearchResult, Side, TreeIter, TreeNodeIter,
+    UnsafeNode,
+};
+
+fn replace<T, R>(v: &mut T, f: impl FnOnce(T) -> (T, R)) -> R {
+    struct Guard;
+    impl Drop for Guard {
+        fn drop(&mut self) {
+            panic!("replace() panicked");
+        }
+    }
+    let guard = Guard;
+    let value = unsafe { core::ptr::read(v) };
+    let (new_value, ret) = f(value);
+    unsafe { core::ptr::write(v, new_value) };
+    core::mem::forget(guard);
+    ret
+}

crates/rbtree/src/raw_node.rs

@@ -54,6 +54,26 @@ pub enum SearchResult<T> {
     Empty,
 }
 
+impl<N: UnsafeNode> SearchResult<Handle<N>> {
+    fn next_extant(self, inclusive: bool) -> Option<Handle<N>> {
+        match self {
+            SearchResult::FoundAt(handle) if inclusive => Some(handle),
+            SearchResult::FoundAt(handle) | SearchResult::NotFoundAt(handle) => handle.next_of(),
+            SearchResult::Empty => None,
+        }
+    }
+
+    fn next_back_extant(self, inclusive: bool) -> Option<Handle<N>> {
+        match self {
+            SearchResult::FoundAt(handle) if inclusive => Some(handle),
+            SearchResult::FoundAt(handle) | SearchResult::NotFoundAt(handle) => {
+                handle.next_back_of()
+            }
+            SearchResult::Empty => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Color {
     Red,
@@ -197,7 +217,7 @@ impl<N> Clone for Handle<N> {
 impl<N> !Sync for Handle<N> {}
 
 impl<N: UnsafeNode> Handle<N> {
-    fn from_node(node: Option<NonNull<N>>) -> Self {
+    pub fn from_node(node: Option<NonNull<N>>) -> Self {
         let Some(node) = node else {
             return Handle::EmptyRoot;
         };
@@ -218,8 +238,7 @@ impl<N: UnsafeNode> Handle<N> {
         }
     }
 
-    #[expect(dead_code)]
-    fn refresh_from_node(&mut self) {
+    pub fn refresh_from_node(&mut self) {
         *self = Self::from_node(self.node());
     }
 
@@ -396,6 +415,34 @@ impl<N: UnsafeNode> Handle<N> {
         Some(side.map(|_| parent))
     }
 
+    fn remove_from_parent_into_node(self) -> (Option<NonNull<N>>, Option<Self>) {
+        match self {
+            Handle::EmptyRoot => (None, None),
+            Handle::Root(node) => (Some(node), None),
+            Handle::Child {
+                parent,
+                node: LeftOrRight::Left(Some(node)),
+            } => {
+                unsafe { parent.as_ref().set_left(None) };
+                (Some(node), Some(Handle::from_node(Some(parent))))
+            }
+            Handle::Child {
+                parent,
+                node: LeftOrRight::Right(Some(node)),
+            } => {
+                unsafe { parent.as_ref().set_right(None) };
+                (Some(node), Some(Handle::from_node(Some(parent))))
+            }
+            Handle::Child {
+                parent,
+                node: LeftOrRight::Right(None) | LeftOrRight::Left(None),
+            } => {
+                unsafe { parent.as_ref().set_right(None) };
+                (None, Some(Handle::from_node(Some(parent))))
+            }
+        }
+    }
+
     fn parent(&self) -> Option<Self> {
         self.parent_and_side().map(LeftOrRight::into_inner)
     }
@@ -528,6 +575,18 @@ impl<N: UnsafeNode> Handle<N> {
         Some(current)
     }
 
+    /// moves to the least non-nil node in the tree, or `Err(self)`.
+    pub fn into_least_leaf(mut self) -> Self {
+        while let Some(child) = self
+            .left_child_extant()
+            .or_else(|| self.right_child_extant())
+        {
+            self = child;
+        }
+
+        self
+    }
+
     /// Returns the greatest non-nil node in the subtree rooted at `self`, or
     /// `None` if the subtree is empty.
     pub fn maximum_of(&self) -> Option<Self> {
@@ -839,7 +898,7 @@ impl<N: UnsafeNode> RBTree<N> {
     pub fn pop_min(&mut self) -> Option<NonNull<N>> {
         let min = self.root_handle().minimum_of()?;
 
-        self.remove_node(min)
+        self.remove_handle(min)
     }
 
     #[must_use]
@@ -852,11 +911,11 @@ impl<N: UnsafeNode> RBTree<N> {
             return None;
         };
 
-        self.remove_node(z)
+        self.remove_handle(z)
     }
 
     #[must_use]
-    fn remove_node(&mut self, z: Handle<N>) -> Option<NonNull<N>> {
+    pub fn remove_handle(&mut self, z: Handle<N>) -> Option<NonNull<N>> {
         // Y is either Z, the removed node in the case that Z has at most
         // one child, or Y is Z's successor which is guaranteed to have at most one
         // child (the right child).
@@ -1028,6 +1087,74 @@ impl<N: UnsafeNode> RBTree<N> {
             range: TreeRange::full_range(self.root_handle()),
         }
     }
+
+    pub fn range<T, R>(&self, range: R) -> TreeIter<'_, N>
+    where
+        T: Ord + ?Sized,
+        N::Key: core::borrow::Borrow<T> + Ord,
+        R: core::ops::RangeBounds<T>,
+    {
+        use core::ops::Bound;
+        let front = match range.start_bound() {
+            Bound::Included(key) => self.find_by_key(key).next_extant(true),
+            Bound::Excluded(key) => self.find_by_key(key).next_extant(false),
+            Bound::Unbounded => None,
+        }
+        .or_else(|| self.root_handle().minimum_of());
+
+        let back = match range.end_bound() {
+            Bound::Included(key) => self.find_by_key(key).next_back_extant(true),
+            Bound::Excluded(key) => self.find_by_key(key).next_back_extant(false),
+            Bound::Unbounded => None,
+        }
+        .or_else(|| self.root_handle().maximum_of());
+
+        TreeIter {
+            range: TreeRange {
+                start: front.map(RangeHandle::Node),
+                end: back.map(RangeHandle::Node),
+                _pd: PhantomData,
+            },
+        }
+    }
+
+    pub fn drain(&mut self) -> Drain<'_, N> {
+        let node = self.root_handle();
+        self.root = None;
+
+        Drain {
+            node,
+            _pd: PhantomData,
+        }
+    }
+}
+
+pub struct Drain<'a, N: UnsafeNode> {
+    node: Handle<N>,
+    _pd: PhantomData<&'a mut RBTree<N>>,
+}
+
+impl<'a, N: UnsafeNode> Drain<'a, N> {
+    fn next(&mut self) -> Option<NonNull<N>> {
+        if self.node.is_nil() {
+            return None;
+        }
+
+        crate::replace(&mut self.node, |node| {
+            let leaf = node.into_least_leaf();
+            let (node, parent) = leaf.remove_from_parent_into_node();
+
+            (parent.unwrap_or(Handle::EmptyRoot), node)
+        })
+    }
+}
+
+impl<'a, N: UnsafeNode> Iterator for Drain<'a, N> {
+    type Item = NonNull<N>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next()
+    }
 }
 
 impl<N: UnsafeNode> Default for RBTree<N> {
@@ -1078,6 +1205,15 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
             _pd: PhantomData,
         }
     }
+
+    fn empty() -> Self {
+        Self {
+            start: None,
+            end: None,
+            _pd: PhantomData,
+        }
+    }
+
     fn full_range(root: Handle<N>) -> Self {
         Self {
             start: Some(RangeHandle::Root(root.clone())),
@@ -1121,8 +1257,8 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
         match self.end {
             None => None,
             _ => {
-                if self.start == self.end {
-                    self.end = None;
+                if current == self.end {
+                    *self = TreeRange::empty();
                 }
 
                 current.map(RangeHandle::into_inner)
@@ -1137,8 +1273,8 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
         match self.start {
             None => None,
             _ => {
-                if self.start == self.end {
-                    self.start = None;
+                if current == self.start {
+                    *self = TreeRange::empty();
                 }
 
                 current.map(RangeHandle::into_inner)
@@ -1195,6 +1331,8 @@ impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeIter<'a, N> {
 mod tests {
     use std::cell::Cell;
 
+    use core::sync::atomic::{AtomicI32, Ordering};
+
     use super::*;
 
     #[derive(Debug)]
@@ -1438,4 +1576,69 @@ mod tests {
             assert!(matches!(result, SearchResult::Empty));
         }
     }
+
+    #[test]
+    fn drain() {
+        let mut tree = RBTree::<TestNode>::new();
+
+        let count = AtomicI32::new(0);
+
+        for node in (1..=10).map(|i| {
+            count.fetch_add(i, Ordering::SeqCst);
+            Box::into_raw(Box::new(TestNode::new(i)))
+        }) {
+            let None = tree.insert_node(unsafe { NonNull::new_unchecked(node) }) else {
+                panic!("duplicate node")
+            };
+        }
+
+        let drain = tree.drain();
+
+        drain.for_each(|node| {
+            let node = unsafe { Box::from_raw(node.as_ptr()) };
+            count.fetch_sub(node.key, Ordering::SeqCst);
+        });
+
+        assert_eq!(
+            count.load(Ordering::SeqCst),
+            0,
+            "all nodes should have been drained and dropped"
+        );
+    }
+
+    #[test]
+    fn drain_collect() {
+        let mut tree = RBTree::<TestNode>::new();
+
+        for node in (1..=10).map(|i| Box::into_raw(Box::new(TestNode::new(i)))) {
+            let None = tree.insert_node(unsafe { NonNull::new_unchecked(node) }) else {
+                panic!("duplicate node")
+            };
+        }
+
+        let mut nodes: Vec<_> = tree
+            .drain()
+            .map(|node| unsafe { Box::from_raw(node.as_ptr()).key })
+            .collect();
+
+        nodes.sort();
+
+        assert_eq!(nodes, (1..=10).collect::<Vec<_>>());
+    }
+
+    #[test]
+    fn range() {
+        let mut tree = RBTree::<TestNode>::new();
+
+        for node in (1..=10).map(|i| Box::into_raw(Box::new(TestNode::new(i)))) {
+            let None = tree.insert_node(unsafe { NonNull::new_unchecked(node) }) else {
+                panic!("duplicate node")
+            };
+        }
+
+        let range = tree.range(3..=7);
+
+        let keys: Vec<_> = range.cloned().collect();
+        assert_eq!(keys, vec![3, 4, 5, 6, 7]);
+    }
 }

kernel/Cargo.toml

@@ -31,3 +31,5 @@ seq-macro = "0.3.6"
 rbtree = { path = "../crates/rbtree" }
 plain = "0.2.3"
 rustc-demangle = "0.1.28"
+num-traits = { version = "0.2.19", default-features = false }
+foundation = { version = "0.1.0", path = "../crates/foundation" }

kernel/src/lib.rs

@@ -8,6 +8,7 @@
     ptr_cast_slice,
     likely_unlikely,
     int_roundings,
+    slice_ptr_get,
     never_type
 )]
 #![cfg_attr(test, feature(custom_test_frameworks))]

kernel/src/main.rs

@@ -27,15 +27,15 @@ extern "C" fn _start() -> ! {
 
     kernel::limine::init_limine_boot_info();
 
-    kernel::serial_println!(
-        "HHDM offset: 0x{:#x}",
-        kernel::memory::HHDM_BASE.get().unwrap()
-    );
+    // kernel::serial_println!(
+    //     "HHDM offset: 0x{:#x}",
+    //     kernel::memory::HHDM_BASE.get().unwrap()
+    // );
 
-    kernel::serial_println!(
-        "Memory map: {:#?}",
-        kernel::limine::MEMMAP_REQUEST.entries()
-    );
+    // kernel::serial_println!(
+    //     "Memory map: {:#?}",
+    //     kernel::limine::MEMMAP_REQUEST.entries()
+    // );
 
     GDT.load();
     IDT.load();

kernel/src/memory.rs

@@ -4,10 +4,11 @@ use core::{
     fmt::Debug,
     num::{NonZero, NonZeroUsize},
     ptr::NonNull,
+    range::Range,
 };
 
 use bit_field::BitField;
-use rbtree::RBTree;
+use rbtree::{RBTree, UnsafeNode};
 
 use crate::{serial_println, sync::OnceLock, x86_64::PAGE_SIZE};
 
@@ -22,7 +23,7 @@ pub trait VirtAddrTranslationExt: Sized {
     }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq)]
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub struct PhyAddr(pub u64);
 
 impl PhyAddr {
@@ -119,11 +120,6 @@ impl VirtAddr {
     }
 }
 
-pub struct PageChunk {
-    pub next: PhyAddr,
-    pub count: usize,
-}
-
 pub struct PhysicalMemoryManager {
     /// On amd64 platforms, there are at most 2^40 pages of addressable physical
     /// memory. We maintain a binary tree of free page chunks for each
@@ -165,20 +161,21 @@ impl Debug for PhysicalMemoryManager {
 
 impl PhysicalMemoryManager {
     pub fn from_memory_map(memory_map: &[crate::boot::MemoryRegion]) -> Self {
-        let usable_regions = memory_map
-            .iter()
-            .filter(|region| region.region_type.is_usable());
+        let mut pmm = PhysicalMemoryManager {
+            buddies: [(); 40].map(|_| RBTree::default()),
+        };
 
-        let buddies = [(); 40].map(|_| RBTree::default());
+        let mut free_usable_regions = free_usable_regions_from_memory_map(memory_map);
 
-        let mut pmm = PhysicalMemoryManager { buddies };
+        serial_println!("Freeing usable regions: {:#?}", free_usable_regions);
 
-        usable_regions.for_each(|region| {
-            let idx = region.start / crate::x86_64::PAGE_SIZE as u64;
-            let count = region.length / crate::x86_64::PAGE_SIZE as u64;
-
-            pmm.free_region(idx as usize, count as usize);
-        });
+        for region in free_usable_regions
+            .drain()
+            .map(|p| unsafe { p.cast::<PageHeader>().as_ref() })
+        {
+            let (idx, count) = region.index_and_count();
+            pmm.free_region(idx, count);
+        }
 
         pmm
     }
@@ -226,6 +223,8 @@ impl PhysicalMemoryManager {
         unsafe { chunk.as_ref().phy().map(|phy| (phy, 1 << bin)) }
     }
 
+    // TODO: grow and shrink
+
     fn free_region(&mut self, page_index: usize, mut count: usize) {
         serial_println!(
             "Freeing region: page_index = {:#x}, count = {:#x}",
@@ -286,7 +285,7 @@ impl PhysicalMemoryManager {
             );
             let None = self.buddies[bin].insert_node(node_ptr) else {
                 panic!(
-                    "Attempted to free a region that is already free: page_index = {}, count = {}",
+                    "Attempted to free a region that is already free: page_index = {:#x}, count = {:#x}",
                     page_index, bit
                 )
             };
@@ -294,10 +293,6 @@ impl PhysicalMemoryManager {
     }
 }
 
-pub struct PhysicalMemoryAllocator {
-    tree: RBTree<PhysicalPageNode>,
-}
-
 #[repr(C)]
 #[derive(Debug)]
 struct PageHeader {
@@ -305,11 +300,79 @@ struct PageHeader {
     count: u64,
 }
 
+unsafe impl UnsafeNode for PageHeader {
+    type Key = PhysicalPageNodeKey;
+
+    fn left(&self) -> Option<NonNull<Self>> {
+        self.node.left().map(Self::unsafe_from_node)
+    }
+
+    fn right(&self) -> Option<NonNull<Self>> {
+        self.node.right().map(Self::unsafe_from_node)
+    }
+
+    fn parent(&self) -> Option<NonNull<Self>> {
+        self.node.parent().map(Self::unsafe_from_node)
+    }
+
+    fn key(&self) -> &Self::Key {
+        self.node.key()
+    }
+
+    fn color(&self) -> rbtree::Color {
+        self.node.color()
+    }
+
+    fn set_left(&self, left: Option<NonNull<Self>>) {
+        self.node.set_left(left.map(Self::unsafe_into_node))
+    }
+
+    fn set_right(&self, right: Option<NonNull<Self>>) {
+        self.node.set_right(right.map(Self::unsafe_into_node))
+    }
+
+    fn set_parent(&self, parent: Option<NonNull<Self>>) {
+        self.node.set_parent(parent.map(Self::unsafe_into_node))
+    }
+
+    fn set_color(&self, color: rbtree::Color) {
+        self.node.set_color(color)
+    }
+}
+
 impl PageHeader {
     fn phy(&self) -> Option<PhyAddr> {
         PhyAddr::from_hhdm_virt(self)
     }
 
+    fn unsafe_from_node(node: NonNull<PhysicalPageNode>) -> NonNull<Self> {
+        node.cast::<Self>()
+    }
+
+    fn unsafe_into_node(ptr: NonNull<Self>) -> NonNull<PhysicalPageNode> {
+        ptr.cast::<PhysicalPageNode>()
+    }
+
+    fn contains(&self, phy: PhyAddr) -> bool {
+        let start = self.phy().expect("PageHeader is not in HHDM");
+        let end = start.page_add(self.count as usize);
+
+        phy >= start && phy < end
+    }
+
+    fn page_range(&self) -> Range<usize> {
+        let (start, count) = self.index_and_count();
+
+        (start..start + count).into()
+    }
+
+    fn index_and_count(&self) -> (usize, usize) {
+        let start = self.phy().expect("PageHeader is not in HHDM").page_index() as usize;
+        let count = self.count as usize;
+
+        (start, count)
+    }
+
     fn new_from_page_idx_and_count(page_idx: usize, count: usize) -> NonNull<Self> {
         let phy = PhyAddr(page_idx as u64 * PAGE_SIZE as u64);
         let virt = phy.as_hhdm_virt();
@@ -326,72 +389,112 @@ impl PageHeader {
     }
 }
 
-impl Debug for PhysicalMemoryAllocator {
-    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        f.debug_struct("PhysicalMemoryAllocator")
-            .field_with("tree", |f| {
-                let iter = self.tree.iter();
-                write!(f, "[")?;
-                for key in iter {
-                    let header = unsafe { (&raw const *key).cast::<PageHeader>().read_volatile() };
+fn free_usable_regions_from_memory_map(
+    memory_map: &[crate::boot::MemoryRegion],
+) -> RBTree<PageHeader> {
+    let mut usable_regions = memory_map
+        .iter()
+        .copied()
+        .filter(|region| region.region_type.is_usable())
+        .fold(RBTree::new(), |mut tree, region| {
+            let idx = region.start as usize / crate::x86_64::PAGE_SIZE;
+            let count = region.length as usize / crate::x86_64::PAGE_SIZE;
 
-                    if let Some(phy) = header.phy() {
-                        writeln!(
-                            f,
-                            "({:?}..{:?}), ",
-                            phy,
-                            phy.page_add(header.count as usize)
-                        )?;
-                    }
-                }
-                write!(f, "]")
-            })
-            .finish()
-    }
-}
+            let node_ptr = PageHeader::new_from_page_idx_and_count(idx, count);
+
+            let None = tree.insert_node(node_ptr) else {
+                panic!(
+                    "Failed to insert node for page_index = {}, count = {}",
+                    idx, count
+                );
+            };
+
+            tree
+        });
 
-impl PhysicalMemoryAllocator {
-    pub fn from_memory_map(memory_map: &[crate::boot::MemoryRegion]) -> Self {
     let mapped_non_usable_regions = memory_map
         .iter()
         .filter(|region| region.region_type.mapped_non_usable());
 
-        let usable_regions = memory_map
-            .iter()
-            .copied()
-            .filter(|region| region.region_type.is_usable());
+    for region in mapped_non_usable_regions {
+        let idx = region.start as usize / crate::x86_64::PAGE_SIZE;
+        let count = region.length as usize / crate::x86_64::PAGE_SIZE;
+        let end = idx + count;
 
-        let mut pmm = PhysicalMemoryAllocator {
-            tree: RBTree::new(),
-        };
+        while let Some(overlapping) = usable_regions
+            .range(
+                PhysicalPageNodeKey::from_page_index(idx)
+                    ..PhysicalPageNodeKey::from_page_index(idx + count),
+            )
+            .map(|p| unsafe {
+                p.phy()
+                    .unwrap()
+                    .as_hhdm_virt()
+                    .as_ptr::<PageHeader>()
+                    .as_ref()
+                    .unwrap()
+            })
+            .find(|header| {
+                header.contains(PhyAddr::from_page_index(idx))
+                    || header.contains(PhyAddr::from_page_index(idx + count - 1))
+            })
+            .inspect(|header| {
+                usable_regions
+                    .remove(header.node.key())
+                    .expect("node exists");
+            })
+        {
+            // We found an overlapping region.
 
-        usable_regions.for_each(|region| {
-            let idx = region.start / crate::x86_64::PAGE_SIZE as u64;
-            let count = region.length / crate::x86_64::PAGE_SIZE as u64;
+            serial_println!("Found overlapping region: {:?}", overlapping);
 
-            pmm.free_region(idx as usize, count as usize);
-        });
+            let Range {
+                start: overlap_idx,
+                end: overlap_end,
+            } = overlapping.page_range();
+            let overlap_count = overlap_end - overlap_idx;
 
-        pmm
-    }
-
-    pub fn free_region(&mut self, page_index: usize, count: usize) {
-        serial_println!(
-            "Freeing region: page_index = {}, count = {}",
-            page_index,
-            count
-        );
-
-        let node_ptr =
-            PageHeader::new_from_page_idx_and_count(page_index, count).cast::<PhysicalPageNode>();
-
-        let None = self.tree.insert_node(node_ptr) else {
-            panic!(
-                "Failed to insert node for page_index = {}, count = {}",
-                page_index, count
-            );
+            // check if our region is a suffix or prefix of the overlapping region
+            if overlap_idx == idx {
+                if let Some(suffix_count) = overlap_count.checked_sub(count) {
+                    let None = usable_regions.insert_node(PageHeader::new_from_page_idx_and_count(
+                        overlap_idx + count,
+                        suffix_count,
+                    )) else {
+                        panic!("pmm: double free")
                     };
                 }
+            } else if overlap_end == end {
+                if let Some(suffix_count) = overlap_count.checked_sub(count) {
+                    let None = usable_regions.insert_node(PageHeader::new_from_page_idx_and_count(
+                        overlap_idx,
+                        suffix_count,
+                    )) else {
+                        panic!("pmm: double free")
+                    };
+                }
+            } else {
+                // region is in the middle of the overlapping region
+                if let Some(prefix_count) = idx.checked_sub(overlap_idx) {
+                    let None = usable_regions.insert_node(PageHeader::new_from_page_idx_and_count(
+                        overlap_idx,
+                        prefix_count,
+                    )) else {
+                        panic!("pmm: double free")
+                    };
+                }
+
+                if let Some(suffix_count) = overlap_end.checked_sub(end) {
+                    let None = usable_regions
+                        .insert_node(PageHeader::new_from_page_idx_and_count(end, suffix_count))
+                    else {
+                        panic!("pmm: double free")
+                    };
+                }
+            }
+        }
+    }
+    usable_regions
 }
 
 #[derive(Debug, Default)]
@@ -629,7 +732,315 @@ unsafe impl rbtree::UnsafeNode for PhysicalPageNode {
     }
 }
 
+pub mod slab {
+    //! A slab allocator
+
+    use core::{
+        alloc::Layout,
+        cell::Cell,
+        hint::{cold_path, unlikely},
+        num::NonZero,
+        ptr::NonNull,
+    };
+
+    const UNLINKED: NonNull<()> = unsafe { NonNull::new_unchecked(!0 as *mut ()) };
+
+    use alloc::alloc::Allocator;
+
+    use crate::x86_64::PAGE_SIZE;
+
+    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>>>,
+        alloc: A,
+    }
+
+    struct SlabChunk<A: Allocator + Clone> {
+        /// Pointer to the next chunk in the slab.
+        next: Option<NonNull<SlabChunk<A>>>,
+        /// Pointer to the slab that owns this chunk.
+        slab: NonNull<Slab<A>>,
+        /// Linked list of free elements in this slab. When this is `None`, the
+        /// slab is full.
+        free: Cell<Option<NonNull<ChunkSlot>>>,
+        /// Number of outstanding allocations from this slab. When this reaches
+        /// zero, the slab can be freed.
+        count: Cell<usize>,
+    }
+
+    struct ChunkSlot(Option<NonNull<Self>>);
+
+    enum SlotResult {
+        Some(NonNull<u8>),
+        Last(NonNull<u8>),
+        None,
+    }
+
+    impl<A: Allocator + Clone> SlabChunk<A> {
+        fn pop_free_slot(&self) -> SlotResult {
+            let Some(slot) = self.free.get() else {
+                return SlotResult::None;
+            };
+
+            self.free.set(unsafe { slot.as_ref() }.next());
+            self.count.update(|count| count + 1);
+
+            match self.free.get() {
+                Some(_) => SlotResult::Some(slot.cast()),
+                None => SlotResult::Last(slot.cast()),
+            }
+        }
+
+        /// returns `true` if the slab is now empty and can be freed
+        fn push_free_slot(&self, slot: NonNull<u8>) -> bool {
+            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
+        }
+    }
+
+    impl ChunkSlot {
+        fn next(&self) -> Option<NonNull<Self>> {
+            self.0
+        }
+    }
+
+    impl<A: Allocator + Clone> Slab<A> {
+        fn new(element_size: usize, alloc: A) -> Self {
+            assert!(
+                element_size.is_power_of_two(),
+                "element_size must be a power of two"
+            );
+
+            Self {
+                element_size,
+                head: None,
+                alloc,
+            }
+        }
+
+        fn first_slot_offset(&self) -> usize {
+            foundation::mem::align_up(core::mem::size_of::<SlabChunk<A>>(), self.element_size)
+        }
+
+        fn count_and_layout(&self) -> (usize, Layout) {
+            let (count, size, align) = {
+                let one_page_count = (PAGE_SIZE - self.first_slot_offset()) / self.element_size;
+
+                if one_page_count < 3 {
+                    let count = 3;
+                    let size =
+                        (self.first_slot_offset() + count * self.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));
+
+                    (count, size, size)
+                } else {
+                    let count = one_page_count;
+                    let size = PAGE_SIZE;
+                    (count, size, self.element_size)
+                }
+            };
+
+            (count, unsafe {
+                Layout::from_size_align_unchecked(size, align)
+            })
+        }
+
+        fn alloc_chunk(&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 Some(bytes) = self.alloc.allocate(layout).ok() else {
+                panic!()
+            };
+
+            let chunk = bytes.as_non_null_ptr().cast::<SlabChunk<A>>();
+            unsafe {
+                let first_slot = chunk
+                    .as_ptr()
+                    .byte_add(self.first_slot_offset())
+                    .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);
+                    chunk.write(ChunkSlot(Some(NonNull::new_unchecked(next))));
+                }
+                first_slot
+                    .byte_add((count - 1) * self.element_size)
+                    .write(ChunkSlot(None));
+
+                chunk.write(SlabChunk {
+                    next: self.head,
+                    slab: NonNull::from(self),
+                    free: Cell::new(Some(NonNull::new_unchecked(first_slot))),
+                    count: Cell::new(0),
+                });
+            }
+
+            chunk
+        }
+
+        #[cold]
+        fn alloc_chunk_cold(&mut self) -> NonNull<SlabChunk<A>> {
+            self.alloc_chunk()
+        }
+
+        fn alloc_slot(&mut self) -> NonNull<[u8]> {
+            let mut chunk = match self.head {
+                Some(chunk) => chunk,
+                None => {
+                    let chunk = self.alloc_chunk_cold();
+                    self.head = Some(chunk);
+                    chunk
+                }
+            };
+
+            let chunk = unsafe { chunk.as_mut() };
+
+            let ptr = match chunk.pop_free_slot() {
+                SlotResult::Some(non_null) => non_null,
+                SlotResult::Last(non_null) => {
+                    self.head = chunk.next.replace(UNLINKED.cast());
+                    non_null
+                }
+                SlotResult::None => {
+                    panic!("SlabChunk is full, but it is still the head of the slab");
+                }
+            };
+
+            ptr.cast_slice(self.element_size)
+        }
+
+        fn free_slot(&mut self, slot: NonNull<u8>) {
+            let (_, layout) = self.count_and_layout();
+            let mut chunk = slot
+                .map_addr(|addr| unsafe {
+                    NonZero::new_unchecked(foundation::mem::align_down(addr.get(), layout.align()))
+                })
+                .cast::<SlabChunk<A>>();
+
+            let chunk = unsafe { chunk.as_mut() };
+            let linked = chunk.next != Some(UNLINKED.cast());
+
+            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 } {
+                        if next == chunk.into() {
+                            unsafe { head.as_mut().next = chunk.next };
+                            break;
+                        }
+                        head = next;
+                    }
+                }
+
+                unsafe {
+                    self.alloc
+                        .deallocate(NonNull::from_mut(chunk).cast(), layout)
+                };
+            } else if !linked {
+                chunk.next = self.head;
+                self.head = Some(chunk.into());
+            }
+        }
+    }
+
+    const SLAB_ALLOCATOR_BUCKETS: usize = 8;
+    pub struct SlabAllocator<A: Allocator + Clone> {
+        /// Slabs for each power-of-two from 16 bytes to 2048 bytes (inclusive).
+        slabs: [Slab<A>; SLAB_ALLOCATOR_BUCKETS],
+        alloc: A,
+    }
+
+    impl<A: Allocator + Clone> SlabAllocator<A> {
+        pub fn new(alloc: A) -> Self {
+            let slabs = [
+                Slab::new(16, alloc.clone()),
+                Slab::new(32, alloc.clone()),
+                Slab::new(64, alloc.clone()),
+                Slab::new(128, alloc.clone()),
+                Slab::new(256, alloc.clone()),
+                Slab::new(512, alloc.clone()),
+                Slab::new(1024, alloc.clone()),
+                Slab::new(2048, alloc.clone()),
+            ];
+
+            Self { slabs, alloc }
+        }
+
+        fn slab_index_for_size(size: usize) -> Option<usize> {
+            if size == 0 {
+                return None;
+            }
+
+            // the smallest slab is 16 bytes
+            let size = size.max(16);
+
+            // get the index of the slab by calculating the log2 of the size and
+            // subtracting 4 (since 2^4 = 16)
+            let index =
+                (size.next_power_of_two().trailing_zeros() - 16usize.trailing_zeros()) as usize;
+
+            // we have 8 slabs, so the index must be less than 8
+            if index < SLAB_ALLOCATOR_BUCKETS {
+                Some(index)
+            } else {
+                None
+            }
+        }
+
+        pub fn alloc(&mut self, layout: Layout) -> Option<NonNull<[u8]>> {
+            if unlikely(layout.size() == 0) {
+                return Some(NonNull::dangling().cast_slice(0));
+            }
+
+            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() })
+                }
+                None => {
+                    // allocate directly from the backing allocator
+                    self.alloc.allocate(layout).ok()
+                }
+            }
+        }
+
+        pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
+            if unlikely(layout.size() == 0) {
+                return;
+            }
+
+            let size = layout.size().max(layout.align());
+
+            match Self::slab_index_for_size(size) {
+                Some(slab_index) => unsafe {
+                    self.slabs.get_unchecked_mut(slab_index).free_slot(ptr)
+                },
+                None => {
+                    // deallocate directly to the backing allocator
+                    unsafe { self.alloc.deallocate(ptr, layout) }
+                }
+            }
+        }
+    }
+}
+
 pub mod bump {
+    //! A bump allocator inspired by / taken from the `stumpalo` crate
+
     use core::{
         alloc::{Allocator, Layout},
         cell::Cell,