kernel: paging part idk..

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!()
+    }
+}