mem, buddy: rework the buddy system

Introduce `Zone`s: a Zone is a region of physical memory that all share
the same NUMA node. The Zone will hold all RawPage structs. Buddy
allocator will now store a reference to the zone that hold all its
pages. Thus, we make the buddy allocator independent of underlying
physical page frame management framework.

Remove unnecessary page flags and structs.

Signed-off-by: greatbridf <greatbridf@icloud.com>

Cargo.lock

@@ -70,7 +70,6 @@ name = "buddy_allocator"
 version = "0.1.0"
 dependencies = [
  "eonix_mm",
- "intrusive_list",
 ]
 
 [[package]]

crates/buddy_allocator/Cargo.toml

@@ -5,4 +5,3 @@ edition = "2024"
 
 [dependencies]
 eonix_mm = { path = "../eonix_mm" }
-intrusive_list = { path = "../intrusive_list" }

crates/buddy_allocator/src/free_area.rs

@@ -1,59 +0,0 @@
-use crate::BuddyRawPage;
-use core::marker::{PhantomData, Send, Sync};
-use intrusive_list::Link;
-
-pub struct FreeArea<T> {
-    free_list: Link,
-    count: usize,
-    _phantom: PhantomData<T>,
-}
-
-unsafe impl<T> Send for FreeArea<T> {}
-unsafe impl<T> Sync for FreeArea<T> {}
-
-impl<Raw> FreeArea<Raw>
-where
-    Raw: BuddyRawPage,
-{
-    pub const fn new() -> Self {
-        Self {
-            free_list: Link::new(),
-            count: 0,
-            _phantom: PhantomData,
-        }
-    }
-
-    pub fn get_free_pages(&mut self) -> Option<Raw> {
-        self.free_list.next_mut().map(|pages_link| {
-            assert_ne!(self.count, 0);
-
-            let pages_ptr = unsafe {
-                // SAFETY: Items in `self.free_list` are guaranteed to be of type `Raw`.
-                Raw::from_link(pages_link)
-            };
-
-            self.count -= 1;
-            pages_link.remove();
-
-            pages_ptr
-        })
-    }
-
-    pub fn add_pages(&mut self, pages_ptr: Raw) {
-        self.count += 1;
-        pages_ptr.set_free();
-
-        unsafe {
-            self.free_list.insert(pages_ptr.get_link());
-        }
-    }
-
-    pub fn del_pages(&mut self, pages_ptr: Raw) {
-        assert!(self.count >= 1 && pages_ptr.is_free());
-        self.count -= 1;
-        pages_ptr.clear_free();
-        unsafe {
-            pages_ptr.get_link().remove();
-        }
-    }
-}

crates/buddy_allocator/src/lib.rs

@@ -1,87 +1,250 @@
 #![no_std]
 
-mod free_area;
-mod zone;
+use core::hint::unreachable_unchecked;
 
-use core::sync::atomic::Ordering;
-use eonix_mm::{
-    address::PAddr,
-    paging::{RawPage, PFN},
-};
-use intrusive_list::Link;
-use zone::Zone;
+use eonix_mm::address::{AddrOps as _, PAddr, PRange};
+use eonix_mm::paging::{PageList, PageListSized, Zone, PFN};
 
 const MAX_ORDER: u32 = 10;
-const ZONE_AREAS: usize = const { MAX_ORDER as usize + 1 };
+const AREAS: usize = const { MAX_ORDER as usize + 1 };
 
-pub trait BuddyRawPage: RawPage {
-    /// Get the container raw page struct of the list link.
-    ///
-    /// # Safety
-    /// The caller MUST ensure that the link points to a `RawPage`.
-    unsafe fn from_link(link: &mut Link) -> Self;
-
-    /// Get the list link of the raw page.
-    ///
-    /// # Safety
-    /// The caller MUST ensure that at any time, only one mutable reference
-    /// to the link exists.
-    unsafe fn get_link(&self) -> &mut Link;
-
-    fn set_order(&self, order: u32);
+pub trait BuddyPage: Sized + 'static {
+    fn pfn(&self) -> PFN;
 
+    fn get_order(&self) -> u32;
     fn is_buddy(&self) -> bool;
-    fn is_free(&self) -> bool;
 
-    fn set_buddy(&self);
-    fn set_free(&self);
+    fn set_order(&mut self, order: u32);
+    fn set_buddy(&mut self, value: bool);
+}
 
-    fn clear_buddy(&self);
-    fn clear_free(&self);
+struct FreeArea<L>
+where
+    L: PageList,
+{
+    free_list: L,
+    count: usize,
 }
 
-pub struct BuddyAllocator<T>
+unsafe impl<L> Send for FreeArea<L> where L: PageList {}
+unsafe impl<L> Sync for FreeArea<L> where L: PageList {}
+
+pub struct BuddyAllocator<Z, L>
 where
-    T: BuddyRawPage,
+    Z: Zone + 'static,
+    L: PageList,
 {
-    zone: Zone<T, ZONE_AREAS>,
+    zone: &'static Z,
+    free_areas: [FreeArea<L>; AREAS],
 }
 
-impl<T> BuddyAllocator<T>
+impl<Z, L> BuddyAllocator<Z, L>
 where
-    T: BuddyRawPage,
+    Z: Zone + 'static,
+    Z::Page: BuddyPage,
+    L: PageListSized,
 {
-    pub const fn new() -> Self {
-        Self { zone: Zone::new() }
+    pub const fn new(zone: &'static Z) -> Self {
+        Self {
+            zone,
+            free_areas: [const { FreeArea::new() }; AREAS],
+        }
     }
+}
 
+impl<Z, L, P> BuddyAllocator<Z, L>
+where
+    Z: Zone<Page = P>,
+    L: PageList<Page = P>,
+    P: BuddyPage + 'static,
+{
     pub fn create_pages(&mut self, start: PAddr, end: PAddr) {
-        self.zone.create_pages(start, end);
+        assert!(
+            self.zone
+                .contains_prange(PRange::new(start.ceil(), end.floor())),
+            "The given address range is not within the zone."
+        );
+
+        let mut pfn = PFN::from(start.ceil());
+        let end_pfn = PFN::from(end.floor());
+
+        while pfn < end_pfn {
+            let mut order = usize::from(pfn).trailing_zeros().min(MAX_ORDER);
+            let new_end_pfn = loop {
+                let new_end = pfn + (1 << order);
+
+                if new_end <= end_pfn {
+                    break new_end;
+                }
+
+                order -= 1;
+            };
+
+            unsafe {
+                // SAFETY: We've checked that the range is within the zone above.
+                self.add_page_unchecked(pfn, order)
+            };
+
+            pfn = new_end_pfn;
+        }
+    }
+
+    fn add_page(&mut self, pfn: PFN, order: u32) {
+        let prange = PRange::from(PAddr::from(pfn)).grow(1 << (order + 12));
+        assert!(
+            self.zone.contains_prange(prange),
+            "The given page is not within the zone."
+        );
+
+        unsafe {
+            // SAFETY: Checks above.
+            self.add_page_unchecked(pfn, order);
+        }
+    }
+
+    unsafe fn add_page_unchecked(&mut self, pfn: PFN, order: u32) {
+        let Some(page) = self.zone.get_page(pfn) else {
+            unsafe { unreachable_unchecked() }
+        };
+
+        unsafe {
+            // SAFETY: The caller ensures that the page is unused.
+            let page_mut = &mut *page.get();
+            self.free_areas[order as usize].add_page(page_mut, order);
+        }
+    }
+
+    fn break_page(&mut self, page: &mut P, order: u32, target_order: u32) {
+        let pfn = page.pfn();
+
+        for order in (target_order..order).rev() {
+            let buddy_pfn = pfn + (1 << order);
+
+            unsafe {
+                // SAFETY: We got the page from `self.free_areas`. Checks are
+                //         done when we've put the page into the buddy system.
+                self.add_page_unchecked(buddy_pfn, order);
+            }
+        }
+
+        page.set_order(target_order);
+    }
+
+    pub fn alloc_order(&mut self, order: u32) -> Option<&'static mut Z::Page> {
+        for current_order in order..AREAS as u32 {
+            let Some(page) = self.free_areas[current_order as usize].get_free_page() else {
+                continue;
+            };
+
+            if current_order > order {
+                self.break_page(page, current_order, order);
+            }
+
+            return Some(page);
+        }
+
+        None
+    }
+
+    pub unsafe fn dealloc(&mut self, page: &'static mut Z::Page) {
+        let mut pfn = page.pfn();
+        let mut order = page.get_order();
+
+        assert!(
+            !page.is_buddy(),
+            "Trying to free a page that is already in the buddy system: {pfn:?}",
+        );
+
+        while order < MAX_ORDER {
+            let buddy_pfn = pfn.buddy_pfn(order);
+            let Some(buddy_page) = self.try_get_buddy(buddy_pfn, order) else {
+                break;
+            };
+
+            self.free_areas[order as usize].remove_page(buddy_page);
+            pfn = pfn.combined_pfn(buddy_pfn);
+            order += 1;
+        }
+
+        self.add_page(pfn, order);
     }
 
-    pub fn alloc_order(&mut self, order: u32) -> Option<T> {
-        let pages_ptr = self.zone.get_free_pages(order);
+    /// This function checks whether the given page is within our [`Zone`] and
+    /// is a free buddy page with the specified order.
+    ///
+    /// We can assure exclusive access to a buddy page of [`order`] if
+    /// - the buddy is within the same [`Zone`] as us.
+    /// - the buddy is a free buddy (in some [`FreeArea`])
+    /// - the buddy has order [`order`]
+    fn try_get_buddy<'a>(&mut self, buddy_pfn: PFN, order: u32) -> Option<&'a mut P> {
+        let buddy_page = self.zone.get_page(buddy_pfn)?;
+
+        unsafe {
+            // SAFETY: We just test whether the page is a buddy.
+            let buddy_page_ref = &*buddy_page.get();
 
-        if let Some(pages_ptr) = pages_ptr {
-            // SAFETY: Memory order here can be Relaxed is for the same reason as that
-            // in the copy constructor of `std::shared_ptr`.
-            pages_ptr.refcount().fetch_add(1, Ordering::Relaxed);
-            pages_ptr.clear_free();
+            if !buddy_page_ref.is_buddy() {
+                return None;
+            }
+
+            // Sad...
+            if buddy_page_ref.get_order() != order {
+                return None;
+            }
+
+            // SAFETY: We have the mutable reference to the buddy allocator.
+            //         So all the pages within are exclusively accessible to us.
+            Some(&mut *buddy_page.get())
+        }
+    }
+}
+
+impl<L> FreeArea<L>
+where
+    L: PageListSized,
+{
+    const fn new() -> Self {
+        Self {
+            free_list: L::NEW,
+            count: 0,
         }
+    }
+}
+
+impl<L> FreeArea<L>
+where
+    L: PageList,
+    L::Page: BuddyPage + 'static,
+{
+    pub fn get_free_page(&mut self) -> Option<&'static mut L::Page> {
+        self.free_list.pop_head().map(|page| {
+            assert_ne!(self.count, 0, "Oops");
+
+            page.set_buddy(false);
+            self.count -= 1;
 
-        pages_ptr
+            page
+        })
     }
 
-    pub unsafe fn dealloc(&mut self, page_ptr: T) {
-        self.zone.free_pages(page_ptr);
+    pub fn add_page(&mut self, page: &'static mut L::Page, order: u32) {
+        page.set_order(order);
+        page.set_buddy(true);
+
+        self.count += 1;
+        self.free_list.push_tail(page);
     }
 
-    pub fn has_management_over(page_ptr: T) -> bool {
-        !page_ptr.is_free() && page_ptr.is_buddy()
+    pub fn remove_page(&mut self, page: &mut L::Page) {
+        assert_ne!(self.count, 0, "Oops");
+        page.set_buddy(false);
+
+        self.count -= 1;
+        self.free_list.remove(page);
     }
 }
 
-pub(self) trait BuddyPFNOps {
+trait BuddyPFNOps {
     fn buddy_pfn(self, order: u32) -> PFN;
     fn combined_pfn(self, buddy_pfn: PFN) -> PFN;
 }

crates/buddy_allocator/src/zone.rs

@@ -1,146 +0,0 @@
-use super::free_area::FreeArea;
-use crate::{BuddyPFNOps as _, BuddyRawPage};
-use core::sync::atomic::Ordering;
-use eonix_mm::{
-    address::{AddrOps as _, PAddr},
-    paging::PFN,
-};
-
-pub(super) struct Zone<T, const AREAS: usize> {
-    free_areas: [FreeArea<T>; AREAS],
-}
-
-impl<Raw, const AREAS: usize> Zone<Raw, AREAS>
-where
-    Raw: BuddyRawPage,
-{
-    pub const fn new() -> Self {
-        Self {
-            free_areas: [const { FreeArea::new() }; AREAS],
-        }
-    }
-
-    pub fn get_free_pages(&mut self, order: u32) -> Option<Raw> {
-        for current_order in order..AREAS as u32 {
-            let pages_ptr = self.free_areas[current_order as usize].get_free_pages();
-            let Some(pages_ptr) = pages_ptr else { continue };
-
-            pages_ptr.set_order(order);
-
-            if current_order > order {
-                self.expand(pages_ptr, current_order, order);
-            }
-
-            assert!(
-                pages_ptr.is_present(),
-                "Page {:?} is not present",
-                pages_ptr.into(),
-            );
-
-            assert!(
-                pages_ptr.is_free(),
-                "Page {:?} is not free",
-                pages_ptr.into(),
-            );
-
-            return Some(pages_ptr);
-        }
-        None
-    }
-
-    fn expand(&mut self, pages_ptr: Raw, order: u32, target_order: u32) {
-        let mut offset = 1 << order;
-        let pages_pfn = Into::<PFN>::into(pages_ptr);
-
-        for order in (target_order..order).rev() {
-            offset >>= 1;
-
-            let split_pages_ptr = Raw::from(pages_pfn + offset);
-            split_pages_ptr.set_order(order);
-            split_pages_ptr.set_buddy();
-            self.free_areas[order as usize].add_pages(split_pages_ptr);
-        }
-    }
-
-    pub fn free_pages(&mut self, mut pages_ptr: Raw) {
-        assert_eq!(pages_ptr.refcount().load(Ordering::Relaxed), 0);
-
-        let mut pfn = Into::<PFN>::into(pages_ptr);
-        let mut current_order = pages_ptr.order();
-
-        assert!(
-            pages_ptr.is_present(),
-            "Freeing a page that is not present: {:?}",
-            pages_ptr.into(),
-        );
-
-        assert!(
-            !pages_ptr.is_free(),
-            "Freeing a page that is free: {:?}",
-            pages_ptr.into(),
-        );
-
-        while current_order < (AREAS - 1) as u32 {
-            let buddy_pfn = pfn.buddy_pfn(current_order);
-            let buddy_pages_ptr = Raw::from(buddy_pfn);
-
-            if !self.buddy_check(buddy_pages_ptr, current_order) {
-                break;
-            }
-
-            pages_ptr.clear_buddy();
-            buddy_pages_ptr.clear_buddy();
-            self.free_areas[current_order as usize].del_pages(buddy_pages_ptr);
-
-            pages_ptr = Raw::from(pfn.combined_pfn(buddy_pfn));
-            pfn = pfn.combined_pfn(buddy_pfn);
-
-            pages_ptr.set_buddy();
-            current_order += 1;
-        }
-
-        pages_ptr.set_order(current_order);
-        self.free_areas[current_order as usize].add_pages(pages_ptr);
-    }
-
-    /// This function checks whether a page is free && is a buddy
-    /// we can coalesce a page and its buddy if
-    /// - the buddy is valid(present) &&
-    /// - the buddy is right now in free_areas &&
-    /// - a page and its buddy have the same order &&
-    /// - a page and its buddy are in the same zone (on smp systems).
-    fn buddy_check(&self, pages_ptr: Raw, order: u32) -> bool {
-        if !pages_ptr.is_present() {
-            return false;
-        }
-        if !pages_ptr.is_free() {
-            return false;
-        }
-        if pages_ptr.order() != order {
-            return false;
-        }
-
-        assert_eq!(pages_ptr.refcount().load(Ordering::Relaxed), 0);
-        true
-    }
-
-    /// Only used on buddy initialization
-    pub fn create_pages(&mut self, start: PAddr, end: PAddr) {
-        let mut start_pfn = PFN::from(start.ceil());
-        let end_pfn = PFN::from(end.floor());
-
-        while start_pfn < end_pfn {
-            let mut order = usize::from(start_pfn)
-                .trailing_zeros()
-                .min((AREAS - 1) as u32);
-
-            while start_pfn + (1 << order) as usize > end_pfn {
-                order -= 1;
-            }
-            let page_ptr = Raw::from(start_pfn);
-            page_ptr.set_buddy();
-            self.free_areas[order as usize].add_pages(page_ptr);
-            start_pfn = start_pfn + (1 << order) as usize;
-        }
-    }
-}

crates/eonix_mm/src/paging.rs

@@ -1,9 +1,13 @@
+mod list;
 mod page;
 mod page_alloc;
 mod pfn;
 mod raw_page;
+mod zone;
 
+pub use list::{PageList, PageListSized};
 pub use page::{Page, PageAccess, PageBlock, PAGE_SIZE, PAGE_SIZE_BITS};
 pub use page_alloc::{GlobalPageAlloc, NoAlloc, PageAlloc};
 pub use pfn::PFN;
 pub use raw_page::{RawPage, UnmanagedRawPage};
+pub use zone::Zone;

crates/eonix_mm/src/paging/list.rs

@@ -0,0 +1,19 @@
+pub trait PageList {
+    type Page;
+
+    fn is_empty(&self) -> bool;
+
+    fn peek_head(&mut self) -> Option<&mut Self::Page>;
+
+    fn pop_head(&mut self) -> Option<&'static mut Self::Page>;
+    fn push_tail(&mut self, page: &'static mut Self::Page);
+    fn remove(&mut self, page: &mut Self::Page);
+}
+
+pub trait PageListSized: PageList + Sized {
+    const NEW: Self;
+
+    fn new() -> Self {
+        Self::NEW
+    }
+}

crates/eonix_mm/src/paging/raw_page.rs

@@ -1,14 +1,13 @@
-use super::PFN;
 use core::sync::atomic::AtomicUsize;
 
+use super::PFN;
+
 /// A `RawPage` represents a page of memory in the kernel. It is a low-level
 /// representation of a page that is used by the kernel to manage memory.
 #[doc(notable_trait)]
 pub trait RawPage: Clone + Copy + From<PFN> + Into<PFN> {
     fn order(&self) -> u32;
     fn refcount(&self) -> &AtomicUsize;
-
-    fn is_present(&self) -> bool;
 }
 
 #[derive(Clone, Copy)]
@@ -45,8 +44,4 @@ impl RawPage for UnmanagedRawPage {
     fn refcount(&self) -> &AtomicUsize {
         &UNMANAGED_RAW_PAGE_CLONE_COUNT
     }
-
-    fn is_present(&self) -> bool {
-        true
-    }
 }

crates/eonix_mm/src/paging/zone.rs

@@ -0,0 +1,20 @@
+use core::cell::UnsafeCell;
+
+#[allow(unused_imports)]
+use super::{Page, PageAlloc, RawPage, PFN};
+use crate::address::PRange;
+
+/// A [`Zone`] holds a lot of [`Page`]s that share the same NUMA node or
+/// "physical location".
+pub trait Zone: Send + Sync {
+    type Page;
+
+    /// Whether the [`range`] is within this [`Zone`].
+    fn contains_prange(&self, range: PRange) -> bool;
+
+    /// Get the [`RawPage`] that [`pfn`] points to.
+    ///
+    /// # Return
+    /// [`None`] if [`pfn`] is not in this [`Zone`].
+    fn get_page(&self, pfn: PFN) -> Option<&UnsafeCell<Self::Page>>;
+}

crates/slab_allocator/src/lib.rs

@@ -2,6 +2,7 @@
 
 use core::ptr::NonNull;
 
+use eonix_mm::paging::{PageList, PageListSized};
 use eonix_sync::Spin;
 
 #[repr(C)]
@@ -10,19 +11,6 @@ pub union SlabSlot {
     data: u8,
 }
 
-pub trait SlabPageList: Sized {
-    type Page: SlabPage;
-
-    fn new() -> Self;
-    fn is_empty(&self) -> bool;
-
-    fn peek_head(&mut self) -> Option<&mut Self::Page>;
-
-    fn pop_head(&mut self) -> Option<&'static mut Self::Page>;
-    fn push_tail(&mut self, page: &'static mut Self::Page);
-    fn remove(&mut self, page: &mut Self::Page);
-}
-
 pub trait SlabPage: Sized + 'static {
     fn get_data_ptr(&self) -> NonNull<[u8]>;
 
@@ -98,7 +86,7 @@ where
 
 pub trait SlabPageAlloc {
     type Page: SlabPage;
-    type PageList: SlabPageList<Page = Self::Page>;
+    type PageList: PageList<Page = Self::Page>;
 
     /// Allocate a page suitable for slab system use. The page MUST come with
     /// its allocation count 0 and next free slot None.
@@ -110,7 +98,7 @@ pub trait SlabPageAlloc {
 
 pub(crate) struct SlabList<T>
 where
-    T: SlabPageList,
+    T: PageList,
 {
     empty_list: T,
     partial_list: T,
@@ -132,6 +120,7 @@ unsafe impl<P, const COUNT: usize> Sync for SlabAlloc<P, COUNT> where P: SlabPag
 impl<L, const COUNT: usize> SlabAlloc<L, COUNT>
 where
     L: SlabPageAlloc,
+    L::PageList: PageListSized,
 {
     pub fn new_in(alloc: L) -> Self {
         Self {
@@ -159,17 +148,23 @@ where
 
 impl<T> SlabList<T>
 where
-    T: SlabPageList,
+    T: PageListSized,
 {
-    fn new(object_size: usize) -> Self {
+    const fn new(object_size: usize) -> Self {
         Self {
-            empty_list: T::new(),
-            partial_list: T::new(),
-            full_list: T::new(),
+            empty_list: T::NEW,
+            partial_list: T::NEW,
+            full_list: T::NEW,
             object_size,
         }
     }
+}
 
+impl<T> SlabList<T>
+where
+    T: PageList,
+    T::Page: SlabPage,
+{
     fn alloc_from_partial(&mut self) -> NonNull<u8> {
         let head = self.partial_list.peek_head().unwrap();
         let slot = head.alloc_slot().unwrap();

src/kernel/mem/page_alloc.rs

@@ -1,19 +1,25 @@
 mod raw_page;
+mod zones;
 
 use core::sync::atomic::Ordering;
 
-use buddy_allocator::{BuddyAllocator, BuddyRawPage as _};
+use buddy_allocator::BuddyAllocator;
 use eonix_mm::address::{AddrOps as _, PRange};
-use eonix_mm::paging::{GlobalPageAlloc as GlobalPageAllocTrait, PageAlloc, PFN};
+use eonix_mm::paging::{
+    GlobalPageAlloc as GlobalPageAllocTrait, PageAlloc, PageList, PageListSized as _, PFN,
+};
+use eonix_preempt::PreemptGuard;
 use eonix_sync::{NoContext, Spin};
-use intrusive_list::List;
-use raw_page::PageFlags;
+use raw_page::{PageFlags, RawPageList};
 pub use raw_page::{RawPage, RawPagePtr};
+pub use zones::GlobalZone;
 
 const COSTLY_ORDER: u32 = 3;
+const AREAS: usize = COSTLY_ORDER as usize + 1;
 const BATCH_SIZE: u32 = 64;
 
-static BUDDY_ALLOC: Spin<BuddyAllocator<RawPagePtr>> = Spin::new(BuddyAllocator::new());
+static BUDDY_ALLOC: Spin<BuddyAllocator<GlobalZone, RawPageList>> =
+    Spin::new(BuddyAllocator::new(&GlobalZone()));
 
 #[eonix_percpu::define_percpu]
 static PERCPU_PAGE_ALLOC: PerCpuPageAlloc = PerCpuPageAlloc::new();
@@ -26,58 +32,42 @@ pub struct BuddyPageAlloc();
 
 struct PerCpuPageAlloc {
     batch: u32,
-    // TODO: might be used in the future.
-    // high: u32,
-    free_areas: [List; COSTLY_ORDER as usize + 1],
+    free_areas: [RawPageList; AREAS],
+}
+
+pub trait PerCpuPage {
+    fn set_local(&mut self, val: bool);
 }
 
 impl PerCpuPageAlloc {
     const fn new() -> Self {
         Self {
             batch: BATCH_SIZE,
-            // high: 0,
-            free_areas: [const { List::new() }; COSTLY_ORDER as usize + 1],
+            free_areas: [RawPageList::NEW; AREAS],
         }
     }
 
-    fn insert_free_pages(&mut self, pages_ptr: RawPagePtr, order: u32) {
-        let free_area = &mut self.free_areas[order as usize];
-        free_area.insert(unsafe { pages_ptr.get_link() });
-    }
-
-    fn get_free_pages(&mut self, order: u32) -> Option<RawPagePtr> {
-        let free_area = &mut self.free_areas[order as usize];
-        free_area.pop().map(|node| unsafe {
-            // SAFETY: `node` is a valid pointer to a `Link` that is not used by anyone.
-            RawPagePtr::from_link(node)
-        })
-    }
-
-    fn alloc_order(&mut self, order: u32) -> Option<RawPagePtr> {
+    fn alloc_order(&mut self, order: u32) -> Option<&'static mut RawPage> {
         assert!(order <= COSTLY_ORDER);
-        if let Some(pages) = self.get_free_pages(order) {
+        if let Some(pages) = self.free_areas[order as usize].pop_head() {
             return Some(pages);
         }
 
         let batch = self.batch >> order;
         for _ in 0..batch {
-            if let Some(pages_ptr) = BUDDY_ALLOC.lock().alloc_order(order) {
-                pages_ptr.flags().set(PageFlags::LOCAL);
-                self.insert_free_pages(pages_ptr, order);
-            } else {
+            let Some(page) = BUDDY_ALLOC.lock().alloc_order(order) else {
                 break;
             };
+
+            page.set_local(true);
+            self.free_areas[order as usize].push_tail(page);
         }
 
-        self.get_free_pages(order)
+        self.free_areas[order as usize].pop_head()
     }
 
-    fn free_pages(&mut self, pages_ptr: RawPagePtr, order: u32) {
-        assert_eq!(pages_ptr.order(), order);
-        assert_eq!(pages_ptr.refcount().load(Ordering::Relaxed), 0);
-
-        pages_ptr.refcount().store(1, Ordering::Relaxed);
-        self.insert_free_pages(pages_ptr, order);
+    fn free_pages(&mut self, page: &'static mut RawPage, order: u32) {
+        self.free_areas[order as usize].push_tail(page);
     }
 }
 
@@ -87,16 +77,6 @@ impl GlobalPageAlloc {
         BuddyPageAlloc()
     }
 
-    pub fn mark_present(range: PRange) {
-        let mut pfn = PFN::from(range.start().ceil());
-        let end_pfn = PFN::from(range.end().floor());
-
-        while pfn < end_pfn {
-            RawPagePtr::from(pfn).flags().set(PageFlags::PRESENT);
-            pfn = pfn + 1;
-        }
-    }
-
     /// Add the pages in the PAddr range `range` to the global allocator.
     ///
     /// This function is only to be called on system initialization when `eonix_preempt`
@@ -116,34 +96,47 @@ impl PageAlloc for GlobalPageAlloc {
     type RawPage = RawPagePtr;
 
     fn alloc_order(&self, order: u32) -> Option<RawPagePtr> {
-        if order > COSTLY_ORDER {
+        let raw_page = if order > COSTLY_ORDER {
             BUDDY_ALLOC.lock().alloc_order(order)
         } else {
             unsafe {
                 eonix_preempt::disable();
-                let page_ptr = PERCPU_PAGE_ALLOC.as_mut().alloc_order(order);
+                let page = PERCPU_PAGE_ALLOC.as_mut().alloc_order(order);
                 eonix_preempt::enable();
-                page_ptr
+
+                page
             }
-        }
+        };
+
+        raw_page.map(|raw_page| {
+            // SAFETY: Memory order here can be Relaxed is for the same reason
+            //         as that in the copy constructor of `std::shared_ptr`.
+            raw_page.refcount.fetch_add(1, Ordering::Relaxed);
+
+            RawPagePtr::from_ref(raw_page)
+        })
     }
 
     unsafe fn dealloc(&self, page_ptr: RawPagePtr) {
+        assert_eq!(
+            page_ptr.refcount().load(Ordering::Relaxed),
+            0,
+            "Trying to free a page with refcount > 0"
+        );
+
         if page_ptr.order() > COSTLY_ORDER {
-            BUDDY_ALLOC.lock().dealloc(page_ptr);
+            BUDDY_ALLOC.lock().dealloc(page_ptr.as_mut());
         } else {
             let order = page_ptr.order();
+
             unsafe {
-                eonix_preempt::disable();
-                PERCPU_PAGE_ALLOC.as_mut().free_pages(page_ptr, order);
-                eonix_preempt::enable();
+                PreemptGuard::new(PERCPU_PAGE_ALLOC.as_mut()).free_pages(page_ptr.as_mut(), order);
             }
         }
     }
 
     fn has_management_over(&self, page_ptr: RawPagePtr) -> bool {
-        BuddyAllocator::has_management_over(page_ptr)
-            && (page_ptr.order() > COSTLY_ORDER || page_ptr.flags().has(PageFlags::LOCAL))
+        page_ptr.order() > COSTLY_ORDER || page_ptr.flags().has(PageFlags::LOCAL)
     }
 }
 
@@ -157,14 +150,17 @@ impl PageAlloc for BuddyPageAlloc {
     type RawPage = RawPagePtr;
 
     fn alloc_order(&self, order: u32) -> Option<RawPagePtr> {
-        BUDDY_ALLOC.lock().alloc_order(order)
+        BUDDY_ALLOC
+            .lock()
+            .alloc_order(order)
+            .map(|raw_page| RawPagePtr::from_ref(raw_page))
     }
 
     unsafe fn dealloc(&self, page_ptr: RawPagePtr) {
-        BUDDY_ALLOC.lock().dealloc(page_ptr);
+        BUDDY_ALLOC.lock().dealloc(page_ptr.as_mut());
     }
 
-    fn has_management_over(&self, page_ptr: RawPagePtr) -> bool {
-        BuddyAllocator::has_management_over(page_ptr)
+    fn has_management_over(&self, _: RawPagePtr) -> bool {
+        true
     }
 }

src/kernel/mem/page_alloc/raw_page.rs

@@ -1,18 +1,18 @@
 use core::ptr::NonNull;
 use core::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
 
-use buddy_allocator::BuddyRawPage;
+use buddy_allocator::BuddyPage;
 use eonix_hal::mm::ArchPhysAccess;
 use eonix_mm::address::{PAddr, PhysAccess as _};
-use eonix_mm::paging::{PageAlloc, RawPage as RawPageTrait, PFN};
+use eonix_mm::paging::{PageAlloc, PageList, PageListSized, RawPage as RawPageTrait, PFN};
 use intrusive_list::{container_of, Link, List};
-use slab_allocator::{SlabPage, SlabPageAlloc, SlabPageList, SlabSlot};
+use slab_allocator::{SlabPage, SlabPageAlloc, SlabSlot};
 
-use super::GlobalPageAlloc;
+use super::{GlobalPageAlloc, PerCpuPage};
 use crate::kernel::mem::page_cache::PageCacheRawPage;
 use crate::kernel::mem::PhysAccess;
 
-const PAGE_ARRAY: NonNull<RawPage> =
+pub const PAGE_ARRAY: NonNull<RawPage> =
     unsafe { NonNull::new_unchecked(0xffffff8040000000 as *mut _) };
 
 pub struct PageFlags(AtomicU32);
@@ -52,21 +52,23 @@ pub struct RawPage {
     /// This field is only used in buddy system and is protected by the global lock.
     order: u32,
     flags: PageFlags,
-    refcount: AtomicUsize,
+    pub refcount: AtomicUsize,
 
     shared_data: PageData,
 }
 
+// XXX: introduce Folio and remove this.
+unsafe impl Send for RawPage {}
+unsafe impl Sync for RawPage {}
+
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct RawPagePtr(NonNull<RawPage>);
 
 impl PageFlags {
-    pub const PRESENT: u32 = 1 << 0;
     pub const LOCKED: u32 = 1 << 1;
     pub const BUDDY: u32 = 1 << 2;
     pub const SLAB: u32 = 1 << 3;
     pub const DIRTY: u32 = 1 << 4;
-    pub const FREE: u32 = 1 << 5;
     pub const LOCAL: u32 = 1 << 6;
 
     pub fn has(&self, flag: u32) -> bool {
@@ -158,48 +160,31 @@ impl RawPageTrait for RawPagePtr {
     fn refcount(&self) -> &AtomicUsize {
         self.refcount()
     }
-
-    fn is_present(&self) -> bool {
-        self.flags().has(PageFlags::PRESENT)
-    }
 }
 
-impl BuddyRawPage for RawPagePtr {
-    unsafe fn from_link(link: &mut Link) -> Self {
-        let raw_page_ptr = container_of!(link, RawPage, link);
-        Self(raw_page_ptr)
-    }
-
-    fn set_order(&self, order: u32) {
-        self.as_mut().order = order;
+impl BuddyPage for RawPage {
+    fn pfn(&self) -> PFN {
+        PFN::from(RawPagePtr::from_ref(self))
     }
 
-    unsafe fn get_link(&self) -> &mut Link {
-        &mut self.as_mut().link
+    fn get_order(&self) -> u32 {
+        self.order
     }
 
     fn is_buddy(&self) -> bool {
-        self.flags().has(PageFlags::BUDDY)
-    }
-
-    fn is_free(&self) -> bool {
-        self.flags().has(PageFlags::FREE)
-    }
-
-    fn set_buddy(&self) {
-        self.flags().set(PageFlags::BUDDY);
-    }
-
-    fn set_free(&self) {
-        self.flags().set(PageFlags::FREE);
+        self.flags.has(PageFlags::BUDDY)
     }
 
-    fn clear_buddy(&self) {
-        self.flags().clear(PageFlags::BUDDY);
+    fn set_order(&mut self, order: u32) {
+        self.order = order;
     }
 
-    fn clear_free(&self) {
-        self.flags().clear(PageFlags::FREE);
+    fn set_buddy(&mut self, val: bool) {
+        if val {
+            self.flags.set(PageFlags::BUDDY);
+        } else {
+            self.flags.clear(PageFlags::BUDDY)
+        }
     }
 }
 
@@ -284,15 +269,21 @@ impl PageCacheRawPage for RawPagePtr {
     }
 }
 
-pub struct RawSlabPageList(List);
+impl PerCpuPage for RawPage {
+    fn set_local(&mut self, val: bool) {
+        if val {
+            self.flags.set(PageFlags::LOCAL)
+        } else {
+            self.flags.clear(PageFlags::LOCAL)
+        }
+    }
+}
+
+pub struct RawPageList(List);
 
-impl SlabPageList for RawSlabPageList {
+impl PageList for RawPageList {
     type Page = RawPage;
 
-    fn new() -> Self {
-        Self(List::new())
-    }
-
     fn is_empty(&self) -> bool {
         self.0.is_empty()
     }
@@ -324,9 +315,13 @@ impl SlabPageList for RawSlabPageList {
     }
 }
 
+impl PageListSized for RawPageList {
+    const NEW: Self = RawPageList(List::new());
+}
+
 impl SlabPageAlloc for GlobalPageAlloc {
     type Page = RawPage;
-    type PageList = RawSlabPageList;
+    type PageList = RawPageList;
 
     unsafe fn alloc_uninit(&self) -> &'static mut RawPage {
         let raw_page = self.alloc().expect("Out of memory").as_mut();

src/kernel/mem/page_alloc/zones.rs

@@ -0,0 +1,25 @@
+use core::cell::UnsafeCell;
+
+use eonix_mm::address::PRange;
+use eonix_mm::paging::{Zone, PFN};
+
+use super::RawPage;
+use crate::kernel::mem::page_alloc::RawPagePtr;
+
+pub struct GlobalZone();
+
+impl Zone for GlobalZone {
+    type Page = RawPage;
+
+    fn contains_prange(&self, _: PRange) -> bool {
+        true
+    }
+
+    fn get_page(&self, pfn: PFN) -> Option<&UnsafeCell<Self::Page>> {
+        unsafe {
+            // SAFETY: The pointer returned by [`RawPagePtr::as_ptr()`] is valid.
+            //         And so is it wrapped with [`UnsafeCell`]
+            Some(&*(RawPagePtr::from(pfn).as_ptr() as *const UnsafeCell<Self::Page>))
+        }
+    }
+}

src/kernel_init.rs

@@ -1,14 +1,11 @@
+use eonix_hal::bootstrap::BootStrapData;
+use eonix_hal::mm::{ArchMemory, ArchPagingMode, GLOBAL_PAGE_TABLE};
+use eonix_hal::traits::mm::Memory;
+use eonix_mm::address::{Addr as _, AddrOps as _, VAddr, VRange};
+use eonix_mm::page_table::{PageAttribute, PagingMode as _, PTE};
+use eonix_mm::paging::{Page as GenericPage, PAGE_SIZE, PFN};
+
 use crate::kernel::mem::{GlobalPageAlloc, RawPage};
-use eonix_hal::{
-    bootstrap::BootStrapData,
-    mm::{ArchMemory, ArchPagingMode, GLOBAL_PAGE_TABLE},
-    traits::mm::Memory,
-};
-use eonix_mm::{
-    address::{Addr as _, AddrOps as _, VAddr, VRange},
-    page_table::{PageAttribute, PagingMode as _, PTE},
-    paging::{Page as GenericPage, PAGE_SIZE, PFN},
-};
 
 pub fn setup_memory(data: &mut BootStrapData) {
     let addr_max = ArchMemory::present_ram()
@@ -50,9 +47,10 @@ pub fn setup_memory(data: &mut BootStrapData) {
         );
     }
 
-    for range in ArchMemory::present_ram() {
-        GlobalPageAlloc::mark_present(range);
-    }
+    // TODO!!!: Construct the global zone with all present ram.
+    // for range in ArchMemory::present_ram() {
+    //     GlobalPageAlloc::mark_present(range);
+    // }
 
     if let Some(early_alloc) = data.take_alloc() {
         for range in early_alloc.into_iter() {