kernel: paging part idk..

crates/foundation/src/alloc/slab.rs

@@ -6,7 +6,10 @@ use core::{
 
 use liballoc::alloc::Allocator;
 
-use crate::mem;
+use crate::{
+    collections::linked_list::{LinkedList, LinkedListNode},
+    mem,
+};
 
 const PAGE_SIZE: usize = 4096;
 
@@ -14,9 +17,9 @@ pub struct Slab<A: Allocator + Clone> {
     /// Size and alignment of each element in the slab.
     element_size: usize,
     /// Pointer to the first chunk in the slab.
-    head: Option<NonNull<SlabChunk<A>>>,
+    chunks: LinkedList<SlabChunk<A>>,
     /// Pointer to the first chunk in the slab that is full.
-    full_head: Option<NonNull<SlabChunk<A>>>,
+    full_chunks: LinkedList<SlabChunk<A>>,
     alloc: A,
     _pd: core::marker::PhantomPinned,
 }
@@ -25,12 +28,11 @@ impl<A: Allocator + Clone> Debug for Slab<A> {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("Slab")
             .field("element_size", &self.element_size)
-            .field("head", &self.head)
+            .finish_non_exhaustive()
-            .field("full_head", &self.full_head)
-            .finish()
     }
 }
 
+#[derive(Debug)]
 struct SlabChunk<A: Allocator + Clone> {
     /// Pointer to the next chunk in the slab.
     next: Option<NonNull<SlabChunk<A>>>,
@@ -44,6 +46,16 @@ struct SlabChunk<A: Allocator + Clone> {
     count: Cell<usize>,
 }
 
+unsafe impl<A: Allocator + Clone> LinkedListNode for SlabChunk<A> {
+    fn next(this: NonNull<Self>) -> Option<NonNull<Self>> {
+        unsafe { (&raw const (*this.as_ptr()).next).read() }
+    }
+
+    fn set_next(this: NonNull<Self>, next: Option<NonNull<Self>>) {
+        unsafe { (&raw mut (*this.as_ptr()).next).write(next) };
+    }
+}
+
 struct ChunkSlot(Option<NonNull<Self>>);
 
 enum SlotResult {
@@ -109,8 +121,8 @@ impl<A: Allocator + Clone> Slab<A> {
 
         let slab = Self {
             element_size,
-            head: None,
+            chunks: LinkedList::new(),
-            full_head: None,
+            full_chunks: LinkedList::new(),
             alloc: alloc.clone(),
             _pd: core::marker::PhantomPinned,
         };
@@ -155,16 +167,17 @@ impl<A: Allocator + Clone> Slab<A> {
         })
     }
 
-    fn alloc_chunk(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
+    fn alloc_chunk(element_size: usize, alloc: &A, slab: NonNull<Self>) -> NonNull<SlabChunk<A>> {
         // we want to limit chunks to 1 page unless the element size is so
         // large that we can fit fewer than 3 elements in a page.
-        let (count, layout) = Self::count_and_layout(self.element_size);
+        let (count, layout) = Self::count_and_layout(element_size);
 
-        let Some(bytes) = self.alloc.allocate(layout).ok() else {
+        let Some(bytes) = alloc.allocate(layout).ok() else {
             panic!()
         };
 
         let chunk = bytes.as_non_null_ptr().cast::<SlabChunk<A>>();
+        let _ = chunk.as_ptr().expose_provenance();
 
         #[cfg(test)]
         std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})");
@@ -172,22 +185,22 @@ impl<A: Allocator + Clone> Slab<A> {
         unsafe {
             let first_slot = chunk
                 .as_ptr()
-                .byte_add(Self::first_slot_offset(self.element_size))
+                .byte_add(Self::first_slot_offset(element_size))
                 .cast::<ChunkSlot>();
 
             for i in 0..(count - 1) {
-                let chunk = first_slot.byte_add(i * self.element_size);
+                let 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))));
             }
             first_slot
-                .byte_add((count - 1) * self.element_size)
+                .byte_add((count - 1) * element_size)
                 .write(ChunkSlot(None));
 
             chunk.write(SlabChunk {
-                next: self.head,
+                next: None,
                 // SAFETY: we only access `slab` via `SlabChunk::slab_pinned_mut`
-                slab: NonNull::from(Pin::into_inner_unchecked(self)),
+                slab,
                 free: Cell::new(Some(NonNull::new_unchecked(first_slot))),
                 count: Cell::new(0),
             });
@@ -196,36 +209,31 @@ impl<A: Allocator + Clone> Slab<A> {
         chunk
     }
 
-    #[cold]
-    fn alloc_chunk_cold(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
-        self.alloc_chunk()
-    }
-
     fn alloc_slot(mut self: Pin<&mut Self>) -> NonNull<[u8]> {
         #[cfg(test)]
         std::eprintln!("Slab::alloc_slot({self:?})");
 
-        let chunk = match self.head {
+        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::Last(non_null) => {
+                let chunk = chunk.remove();
                 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
@@ -241,11 +249,14 @@ impl<A: Allocator + Clone> Slab<A> {
     fn free_slot(element_size: usize, slot: NonNull<u8>) {
         let (_, layout) = Self::count_and_layout(element_size);
 
-        let chunk_ptr = slot
+        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()))
+        let chunk_ptr: NonNull<SlabChunk<A>> = slot.with_addr(addr).cast();
-            })
+
-            .cast::<SlabChunk<A>>();
+        // grab provenance exposed in alloc_chunk, since `slot` may be noalias restricted to its layout.
+
+        // let chunk_ptr: NonNull<SlabChunk<A>> = NonNull::with_exposed_provenance(addr);
+
 
         #[cfg(test)]
         std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})");
@@ -258,47 +269,35 @@ impl<A: Allocator + Clone> Slab<A> {
             FreeSlotResult::Empty => {
                 // chunk is empty:
                 // first unlink the chunk from the slab..
-                if slab.head == Some(chunk.into()) {
+                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 => {
                 // chunk was full:
 
                 // 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 => {
                 // chunk is not empty, and was not full, so nothing to do
@@ -311,26 +310,21 @@ impl<A: Allocator + Clone> Drop for Slab<A> {
     fn drop(&mut self) {
         let (_, layout) = Self::count_and_layout(self.element_size);
 
-        let mut chunk = self.head.take();
+        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)]
-            std::eprintln!("Slab::drop({chunk_ptr:?}, layout: {layout:?})");
+            std::eprintln!("Slab::drop({chunk:?}, layout: {layout:?})");
 
             unsafe {
-                self.alloc.deallocate(chunk_ptr.cast(), layout);
+                self.alloc.deallocate(chunk.cast(), layout);
             }
         }
 
-        let mut chunk = self.full_head.take();
+        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 {
-                self.alloc.deallocate(chunk_ptr.cast(), layout);
+                self.alloc.deallocate(chunk.cast(), layout);
             }
         }
     }
@@ -438,7 +432,7 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
 mod tests {
     use super::*;
 
-    use core::{mem::forget, pin};
+    use core::pin;
     use std::{prelude::rust_2024::*, sync::Mutex};
 
     struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex<Pin<&'a mut SlabAllocator<A>>>);

crates/foundation/src/collections/linked_list.rs

@@ -0,0 +1,216 @@
+use core::alloc::{Allocator, Layout};
+use core::ptr::NonNull;
+
+use crate::mem::DormantMutRef;
+
+pub unsafe trait LinkedListNode {
+    fn next(this: NonNull<Self>) -> Option<NonNull<Self>>;
+    fn set_next(this: NonNull<Self>, next: Option<NonNull<Self>>);
+}
+
+#[derive(Debug)]
+pub struct LinkedList<T: LinkedListNode> {
+    head: Option<NonNull<T>>,
+}
+
+pub struct NextOf<'a, T: LinkedListNode> {
+    parent: NonNull<T>,
+    node: NonNull<T>,
+    _marker: core::marker::PhantomData<&'a mut T>,
+}
+pub struct Head<'a, T: LinkedListNode> {
+    list: DormantMutRef<'a, LinkedList<T>>,
+    node: NonNull<T>,
+}
+
+pub enum LinkedListEntry<'a, T: LinkedListNode> {
+    NextOf(NextOf<'a, T>),
+    Head(Head<'a, T>),
+}
+
+impl<'a, T: LinkedListNode> LinkedListEntry<'a, T> {
+    unsafe fn new_head(list: &'a mut LinkedList<T>, head: NonNull<T>) -> Self {
+        Self::Head(Head {
+            list: DormantMutRef::new(list).1,
+            node: head,
+        })
+    }
+    pub fn get(&self) -> &T {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => unsafe { node.as_ref() },
+            LinkedListEntry::Head(Head { node, .. }) => unsafe { node.as_ref() },
+        }
+    }
+
+    pub fn get_mut(&mut self) -> &mut T {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => unsafe { node.as_mut() },
+            LinkedListEntry::Head(Head { node, .. }) => unsafe { node.as_mut() },
+        }
+    }
+
+    pub fn remove(self) -> NonNull<T> {
+        match self {
+            LinkedListEntry::NextOf(NextOf { parent, node, .. }) => {
+                let next = T::next(node);
+                T::set_next(parent, next);
+                node
+            }
+            LinkedListEntry::Head(Head { list, node }) => {
+                let next = T::next(node);
+                unsafe { list.awaken() }.head = next;
+                node
+            }
+        }
+    }
+
+    pub fn inner(&self) -> NonNull<T> {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => *node,
+            LinkedListEntry::Head(Head { node, .. }) => *node,
+        }
+    }
+    pub fn next(&self) -> Option<NonNull<T>> {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => T::next(*node),
+            LinkedListEntry::Head(Head { node, .. }) => T::next(*node),
+        }
+    }
+}
+
+impl<T: LinkedListNode> LinkedList<T> {
+    pub const fn new() -> Self {
+        Self { head: None }
+    }
+
+    pub fn get_or_insert_front_with<F>(&mut self, f: F) -> LinkedListEntry<'_, T>
+    where
+        F: FnOnce() -> NonNull<T>,
+    {
+        let head = if let Some(head) = self.head {
+            head
+        } else {
+            let node = f();
+            self.push_front(node);
+            node
+        };
+
+        // SAFETY: head is the head of the list
+        unsafe { LinkedListEntry::new_head(self, head) }
+    }
+
+    pub fn push_front(&mut self, node: NonNull<T>) {
+        T::set_next(node, self.head);
+        self.head = Some(node);
+    }
+
+    pub fn pop_front(&mut self) -> Option<NonNull<T>> {
+        let node = self.head?;
+        self.head = T::next(node);
+        Some(node)
+    }
+
+    pub fn remove_if<F>(&mut self, predicate: F) -> RemoveIf<'_, T, F>
+    where
+        F: FnMut(NonNull<T>) -> bool,
+    {
+        RemoveIf::new(self, predicate)
+    }
+
+    pub fn iter(&self) -> LinkedListIter<T> {
+        LinkedListIter { current: self.head }
+    }
+    pub fn into_iter(self) -> LinkedListIter<T> {
+        LinkedListIter { current: self.head }
+    }
+    pub fn entries(&mut self) -> LinkedListEntries<'_, T> {
+        LinkedListEntries::new_head(self)
+    }
+
+    pub unsafe fn drop_in<A: Allocator>(mut self, alloc: &A) {
+        let iter = LinkedListIter {
+            current: self.head.take(),
+        };
+        for node in iter {
+            unsafe { alloc.deallocate(node.cast(), Layout::new::<T>()) }
+        }
+    }
+}
+
+impl<T: LinkedListNode> Default for LinkedList<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pub struct RemoveIf<'a, T: LinkedListNode, F> {
+    entries: LinkedListEntries<'a, T>,
+    pred: F,
+}
+
+impl<'a, T: LinkedListNode, F> RemoveIf<'a, T, F> {
+    pub fn new(list: &'a mut LinkedList<T>, pred: F) -> Self {
+        Self {
+            entries: LinkedListEntries::new_head(list),
+            pred,
+        }
+    }
+}
+
+impl<'a, T: LinkedListNode, F> Iterator for RemoveIf<'a, T, F>
+where
+    F: FnMut(NonNull<T>) -> bool,
+{
+    type Item = NonNull<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(entry) = self.entries.next() {
+            let node = entry.inner();
+            if (self.pred)(node) {
+                return Some(entry.remove());
+            }
+        }
+        None
+    }
+}
+
+pub struct LinkedListEntries<'a, T: LinkedListNode> {
+    current: Option<LinkedListEntry<'a, T>>,
+}
+
+impl<'a, T: LinkedListNode> LinkedListEntries<'a, T> {
+    fn new_head(list: &'a mut LinkedList<T>) -> Self {
+        let current = list
+            .head
+            .map(|head| unsafe { LinkedListEntry::new_head(list, head) });
+        Self { current }
+    }
+    pub fn next(&mut self) -> Option<LinkedListEntry<'_, T>> {
+        if let Some(current) = self.current.take() {
+            self.current = match current.next() {
+                Some(next_node) => Some(LinkedListEntry::NextOf(NextOf {
+                    parent: current.inner(),
+                    node: next_node,
+                    _marker: core::marker::PhantomData,
+                })),
+                None => None,
+            };
+            Some(current)
+        } else {
+            None
+        }
+    }
+}
+
+pub struct LinkedListIter<T: LinkedListNode> {
+    current: Option<NonNull<T>>,
+}
+impl<T: LinkedListNode> Iterator for LinkedListIter<T> {
+    type Item = NonNull<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let node = self.current?;
+        self.current = T::next(node);
+        Some(node)
+    }
+}

crates/foundation/src/collections/mod.rs

@@ -0,0 +1 @@
+pub mod linked_list;

crates/foundation/src/lib.rs

@@ -15,6 +15,7 @@ extern crate alloc as liballoc;
 extern crate std;
 
 pub mod alloc;
+pub mod collections;
 pub mod sync;
 pub mod mem {
     use num_traits::PrimInt;
@@ -42,6 +43,10 @@ pub mod mem {
         value & !(alignment - T::one())
     }
 
+    pub fn is_aligned<T: PrimInt>(value: T, alignment: T) -> bool {
+        value & (alignment - T::one()) == T::zero()
+    }
+
     /// # Safety
     pub unsafe fn volatile_copy<T: Sized>(src: *const T, dst: *mut T, count: usize) {
         unsafe {
@@ -53,6 +58,79 @@ pub mod mem {
             }
         }
     }
+
+    use core::marker::PhantomData;
+    use core::ptr::NonNull;
+
+    /// Models a reborrow of some unique reference, when you know that the reborrow
+    /// and all its descendants (i.e., all pointers and references derived from it)
+    /// will not be used any more at some point, after which you want to use the
+    /// original unique reference again.
+    ///
+    /// The borrow checker usually handles this stacking of borrows for you, but
+    /// some control flows that accomplish this stacking are too complicated for
+    /// the compiler to follow. A `DormantMutRef` allows you to check borrowing
+    /// yourself, while still expressing its stacked nature, and encapsulating
+    /// the raw pointer code needed to do this without undefined behavior.
+    pub(super) struct DormantMutRef<'a, T> {
+        ptr: NonNull<T>,
+        _marker: PhantomData<&'a mut T>,
+    }
+
+    unsafe impl<'a, T> Sync for DormantMutRef<'a, T> where &'a mut T: Sync {}
+    unsafe impl<'a, T> Send for DormantMutRef<'a, T> where &'a mut T: Send {}
+
+    impl<'a, T> DormantMutRef<'a, T> {
+        /// Capture a unique borrow, and immediately reborrow it. For the compiler,
+        /// the lifetime of the new reference is the same as the lifetime of the
+        /// original reference, but you promise to use it for a shorter period.
+        pub(super) fn new(t: &'a mut T) -> (&'a mut T, Self) {
+            let ptr = NonNull::from(t);
+            // SAFETY: we hold the borrow throughout 'a via `_marker`, and we expose
+            // only this reference, so it is unique.
+            let new_ref = unsafe { &mut *ptr.as_ptr() };
+            (
+                new_ref,
+                Self {
+                    ptr,
+                    _marker: PhantomData,
+                },
+            )
+        }
+
+        /// Revert to the unique borrow initially captured.
+        ///
+        /// # Safety
+        ///
+        /// The reborrow must have ended, i.e., the reference returned by `new` and
+        /// all pointers and references derived from it, must not be used anymore.
+        pub(super) unsafe fn awaken(self) -> &'a mut T {
+            // SAFETY: our own safety conditions imply this reference is again unique.
+            unsafe { &mut *self.ptr.as_ptr() }
+        }
+
+        /// Borrows a new mutable reference from the unique borrow initially captured.
+        ///
+        /// # Safety
+        ///
+        /// The reborrow must have ended, i.e., the reference returned by `new` and
+        /// all pointers and references derived from it, must not be used anymore.
+        pub(super) unsafe fn reborrow(&mut self) -> &'a mut T {
+            // SAFETY: our own safety conditions imply this reference is again unique.
+            unsafe { &mut *self.ptr.as_ptr() }
+        }
+
+        /// Borrows a new shared reference from the unique borrow initially captured.
+        ///
+        /// # Safety
+        ///
+        /// The reborrow must have ended, i.e., the reference returned by `new` and
+        /// all pointers and references derived from it, must not be used anymore.
+        pub(super) unsafe fn reborrow_shared(&self) -> &'a T {
+            // SAFETY: our own safety conditions imply this reference is again unique.
+            unsafe { &*self.ptr.as_ptr() }
+        }
+    }
 }
 
 pub struct DropGuard<F: FnOnce()>(::core::mem::ManuallyDrop<F>);

kernel/src/memory.rs

@@ -10,7 +10,11 @@ use core::{
 use bit_field::BitField;
 use rbtree::{RBTree, UnsafeNode};
 
-use crate::{serial_println, sync::OnceLock, x86_64::PAGE_SIZE};
+use crate::{
+    serial_println,
+    sync::{LazyLock, OnceLock, SpinMutex},
+    x86_64::PAGE_SIZE,
+};
 
 pub static HHDM_BASE: OnceLock<u64> = OnceLock::new();
 
@@ -136,6 +140,8 @@ pub struct PhysicalMemoryManager {
     buddies: [RBTree<PhysicalPageNode>; 40],
 }
 
+unsafe impl Send for PhysicalMemoryManager {}
+
 impl Debug for PhysicalMemoryManager {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("PhysicalMemoryManager")
@@ -160,6 +166,17 @@ impl Debug for PhysicalMemoryManager {
 }
 
 impl PhysicalMemoryManager {
+    pub fn get() -> &'static SpinMutex<PhysicalMemoryManager> {
+        static PMM: LazyLock<SpinMutex<PhysicalMemoryManager>> = LazyLock::new(|| {
+            SpinMutex::new(PhysicalMemoryManager::from_memory_map(
+                crate::boot::BOOT_INFO
+                    .get()
+                    .expect("BOOT_INFO is not initialized")
+                    .memory_map,
+            ))
+        });
+        PMM.get().expect("PhysicalMemoryManager is not initialized")
+    }
     pub fn from_memory_map(memory_map: &[crate::boot::MemoryRegion]) -> Self {
         let mut pmm = PhysicalMemoryManager {
             buddies: [(); 40].map(|_| RBTree::default()),
@@ -731,740 +748,3 @@ unsafe impl rbtree::UnsafeNode for PhysicalPageNode {
         self.set_color_bit(color == rbtree::Color::Red);
     }
 }
-
-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,
-        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);
-    }
-}

kernel/src/sync/mod.rs

@@ -430,6 +430,14 @@ mod once {
             }
         }
 
+        pub fn get_mut(&mut self) -> Option<&mut T> {
+            if self.once.is_completed() {
+                Some(unsafe { (&mut *self.t.get()).assume_init_mut() })
+            } else {
+                None
+            }
+        }
+
         fn force(this: &Self) -> &T {
             this.once.call_once(|_| {
                 // SAFETY: because `call_once` will panic if poisoned, this

kernel/src/x86_64/paging.rs

@@ -2,14 +2,17 @@ use core::{
     borrow::Borrow,
     fmt::Debug,
     hint::unlikely,
-    ops::{Deref, Index},
+    marker::PhantomData,
+    ops::{Deref, Index, IndexMut, Range},
 };
 
 use bit_field::BitField;
+use foundation::mem::{align_down, align_up, is_aligned};
 
 use crate::{
+    boot::MemoryRegion,
     memory::{PhyAddr, VirtAddr, VirtAddrTranslationExt},
-    x86_64::{VirtAddrExt, registers::Cr4},
+    x86_64::{PAGE_SIZE, VirtAddrExt, registers::Cr4},
 };
 
 #[repr(C, align(4096))]
@@ -42,6 +45,141 @@ impl Index<u16> for PageTable {
     }
 }
 
+impl IndexMut<u16> for PageTable {
+    fn index_mut(&mut self, index: u16) -> &mut Self::Output {
+        assert!(index < 512, "Page table index out of bounds");
+        &mut self.entries[index as usize]
+    }
+}
+
+#[repr(transparent)]
+struct RootPageTable {
+    table: PageTable,
+}
+
+pub enum PageType {
+    FourKb,
+    TwoMb,
+    OneGb,
+}
+
+impl PageType {
+    const fn page_size(&self) -> u64 {
+        match self {
+            PageType::FourKb => 4 * 1024,
+            PageType::TwoMb => 2 * 1024 * 1024,
+            PageType::OneGb => 1024 * 1024 * 1024,
+        }
+    }
+
+    fn is_4kb(&self) -> bool {
+        matches!(self, PageType::FourKb)
+    }
+    fn is_2mb(&self) -> bool {
+        matches!(self, PageType::TwoMb)
+    }
+    fn is_1gb(&self) -> bool {
+        matches!(self, PageType::OneGb)
+    }
+    fn is_aligned(&self, addr: VirtAddr) -> bool {
+        is_aligned(addr.0, self.page_size())
+    }
+}
+
+enum PageError {
+    AlreadyPresent,
+}
+
+impl RootPageTable {
+    const FOUR_KB: u64 = 4 * 1024;
+    const TWO_MB: u64 = 2 * 1024 * 1024;
+    const ONE_GB: u64 = 1024 * 1024 * 1024;
+    unsafe fn map_to(&mut self, phy: Range<PhyAddr>, virt: VirtAddr, flags: PageFlags) {
+        let start = align_down(phy.start.0, PAGE_SIZE as u64);
+        let end = align_up(phy.end.0, PAGE_SIZE as u64);
+        let size = end - start;
+
+        match size {
+            Self::ONE_GB.. if is_aligned(virt.0, Self::ONE_GB) => {}
+            Self::TWO_MB.. if is_aligned(virt.0, Self::TWO_MB) => {}
+            Self::FOUR_KB.. => {
+                assert!(
+                    is_aligned(virt.0, Self::FOUR_KB),
+                    "Virtual address is not aligned to 4KB"
+                );
+
+                let entry = PageTableEntry::from_addr_and_page_flags(PhyAddr(start), flags);
+            }
+            _ => {
+                unreachable!()
+            }
+        }
+
+        todo!()
+    }
+
+    fn insert(
+        &mut self,
+        virt: VirtAddr,
+        entry: PageTableEntry,
+        page_type: PageType,
+    ) -> Result<(), PageError> {
+        assert!(
+            page_type.is_aligned(virt),
+            "Virtual address is not aligned to page size"
+        );
+        let pml4_index = virt.page_table_index::<{ VirtAddr::PML4 }>();
+        let pdpt_index = virt.page_table_index::<{ VirtAddr::PDPT }>();
+        let pd_index = (!page_type.is_1gb()).then(|| virt.page_table_index::<{ VirtAddr::PD }>());
+        let pt_index = page_type
+            .is_4kb()
+            .then(|| virt.page_table_index::<{ VirtAddr::PT }>());
+
+        let descend = |entry: &mut PageTableEntry, idx: u16| {
+            if entry.contains(PageTableEntryFlags::HUGE_PAGE) {
+                return Err(PageError::AlreadyPresent);
+            }
+            if !entry.present() {
+                let (page, _) = crate::memory::PhysicalMemoryManager::get()
+                    .lock()
+                    .allocate_pages(1)
+                    .expect("Failed to allocate page for page table");
+                entry.set_phy(page);
+                entry.set_present(true);
+            }
+
+            let as_table = unsafe {
+                entry
+                    .phy()
+                    .as_hhdm_virt()
+                    .as_mut::<PageTable>()
+                    .as_mut()
+                    .unwrap_unchecked()
+            };
+
+            Ok(&mut as_table.entries[idx as usize])
+        };
+
+        let mut cursor = &mut self.table[pml4_index];
+        cursor = descend(cursor, pdpt_index)?;
+        if let Some(pd_index) = pd_index {
+            cursor = descend(cursor, pd_index)?;
+
+            if let Some(pt_index) = pt_index {
+                cursor = descend(cursor, pt_index)?;
+            }
+        }
+
+        if cursor.present() {
+            return Err(PageError::AlreadyPresent);
+        } else {
+            *cursor = entry;
+        }
+
+        Ok(())
+    }
+}
+
 #[repr(transparent)]
 #[derive(Clone, Copy, PartialEq, Eq)]
 pub struct PageTableEntry(PageTableEntryFlags);
@@ -81,6 +219,32 @@ impl Deref for PageTableEntry {
     }
 }
 
+bitflags::bitflags! {
+    #[repr(transparent)]
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    pub struct PageFlags: u8 {
+        const WRITE = 1 << 0;
+        const USER = 1 << 1;
+        const EXECUTE = 1 << 2;
+    }
+}
+
+impl PageFlags {
+    pub fn into_pte_flags(self) -> PageTableEntryFlags {
+        let mut flags = PageTableEntryFlags::empty();
+        if self.contains(PageFlags::WRITE) {
+            flags |= PageTableEntryFlags::WRITABLE;
+        }
+        if self.contains(PageFlags::USER) {
+            flags |= PageTableEntryFlags::USER_ACCESSIBLE;
+        }
+        if !self.contains(PageFlags::EXECUTE) {
+            flags |= PageTableEntryFlags::NO_EXECUTE;
+        }
+        flags
+    }
+}
+
 bitflags::bitflags! {
     #[repr(transparent)]
     #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -101,6 +265,18 @@ bitflags::bitflags! {
 }
 
 impl PageTableEntry {
+    pub fn from_addr_and_page_flags(phy: PhyAddr, flags: PageFlags) -> Self {
+        let mut entry = Self(flags.into_pte_flags() | PageTableEntryFlags::PRESENT);
+        entry.as_mut_raw().set_bits(12..52, phy.0 >> 12);
+        entry
+    }
+
+    pub unsafe fn from_addr_and_flags(phy: PhyAddr, flags: PageTableEntryFlags) -> Self {
+        let mut entry = Self(flags);
+        entry.as_mut_raw().set_bits(12..52, phy.0 >> 12);
+        entry
+    }
+
     pub fn from_raw(bits: u64) -> Self {
         Self(PageTableEntryFlags::from_bits_retain(bits))
     }
@@ -113,9 +289,22 @@ impl PageTableEntry {
     pub fn present(&self) -> bool {
         self.contains(PageTableEntryFlags::PRESENT)
     }
+    pub fn set_present(&mut self, present: bool) {
+        if present {
+            self.0 |= PageTableEntryFlags::PRESENT;
+        } else {
+            self.0.remove(PageTableEntryFlags::PRESENT);
+        }
+    }
+
     pub fn phy(&self) -> PhyAddr {
         PhyAddr(self.as_raw().get_bits(12..52) << 12)
     }
+
+    pub fn set_phy(&mut self, phy: PhyAddr) {
+        self.as_mut_raw().set_bits(12..52, phy.0 >> 12);
+    }
+
     pub fn try_as_page_table(&self) -> Option<&PageTable> {
         if !self.contains(PageTableEntryFlags::PRESENT) {
             return None;
@@ -134,6 +323,22 @@ impl PageTableEntry {
             )
         }
     }
+
+    pub unsafe fn as_page_table_mut(&self) -> &mut PageTable {
+        // entry must be present
+        assert!(self.contains(PageTableEntryFlags::PRESENT));
+        // if the entry is a huge page, it does not point to a deeper page table.
+        assert!(!self.contains(PageTableEntryFlags::HUGE_PAGE));
+
+        unsafe {
+            self.phy()
+                .as_hhdm_virt()
+                .as_mut::<PageTable>()
+                .as_mut()
+                .unwrap_unchecked()
+        }
+    }
+
     pub unsafe fn as_page_table(&self) -> &PageTable {
         // entry must be present
         assert!(self.contains(PageTableEntryFlags::PRESENT));
@@ -231,3 +436,42 @@ pub fn get_physical_addr(virt: VirtAddr) -> Option<PhyAddr> {
 
     Some(PhyAddr(phy_addr))
 }
+
+struct Mapping {
+    root: PageTableEntry,
+}
+
+struct MappingBuilder<I> {
+    offset: Option<(VirtAddr, I)>,
+}
+
+impl<I> MappingBuilder<I>
+where
+    I: Iterator<Item = MemoryRegion>,
+{
+    fn with_offset<T>(mut self, offset: VirtAddr, memory_map: T) -> MappingBuilder<T>
+    where
+        T: Iterator<Item = MemoryRegion>,
+    {
+        MappingBuilder {
+            offset: Some((offset, memory_map)),
+        }
+    }
+    fn build(self) -> Mapping {
+        let (root, _) = crate::memory::PhysicalMemoryManager::get()
+            .lock()
+            .allocate_pages(1)
+            .expect("Failed to allocate page for mapping root");
+
+        let root = root.into_hhdm_virt();
+        let root_mut = unsafe {
+            let root_mut = root.as_mut::<PageTable>();
+            root_mut.write(PageTable {
+                entries: [PageTableEntry::from_raw(0); 512],
+            });
+
+            root_mut.as_mut_unchecked()
+        };
+        todo!()
+    }
+}