Allow specifying initial page metadata when allocating

ostd/src/mm/frame/options.rs

@@ -57,9 +57,10 @@ impl FrameAllocOptions {
     /// Allocates a collection of page frames according to the given options.
     pub fn alloc(&self) -> Result<Vec<Frame>> {
         let pages = if self.is_contiguous {
-            page::allocator::alloc(self.nframes * PAGE_SIZE).ok_or(Error::NoMemory)?
+            page::allocator::alloc(self.nframes * PAGE_SIZE, |_| FrameMeta::default())
+                .ok_or(Error::NoMemory)?
         } else {
-            page::allocator::alloc_contiguous(self.nframes * PAGE_SIZE)
+            page::allocator::alloc_contiguous(self.nframes * PAGE_SIZE, |_| FrameMeta::default())
                 .ok_or(Error::NoMemory)?
                 .into()
         };
@@ -79,7 +80,7 @@ impl FrameAllocOptions {
             return Err(Error::InvalidArgs);
         }
 
-        let page = page::allocator::alloc_single().ok_or(Error::NoMemory)?;
+        let page = page::allocator::alloc_single(FrameMeta::default()).ok_or(Error::NoMemory)?;
         let frame = Frame { page };
         if !self.uninit {
             frame.writer().fill(0);
@@ -98,7 +99,7 @@ impl FrameAllocOptions {
         }
 
         let segment: Segment =
-            page::allocator::alloc_contiguous::<FrameMeta>(self.nframes * PAGE_SIZE)
+            page::allocator::alloc_contiguous(self.nframes * PAGE_SIZE, |_| FrameMeta::default())
                 .ok_or(Error::NoMemory)?
                 .into();
         if !self.uninit {

ostd/src/mm/kspace.rs

@@ -211,7 +211,7 @@ pub fn init_kernel_page_table(meta_pages: Vec<Page<MetaPageMeta>>) {
         };
         let mut cursor = kpt.cursor_mut(&from).unwrap();
         for frame_paddr in to.step_by(PAGE_SIZE) {
-            let page = Page::<KernelMeta>::from_unused(frame_paddr);
+            let page = Page::<KernelMeta>::from_unused(frame_paddr, KernelMeta::default());
             // SAFETY: we are doing mappings for the kernel.
             unsafe {
                 cursor.map(page.into(), prop);

ostd/src/mm/page/allocator.rs

@@ -13,7 +13,11 @@ use log::info;
 use spin::Once;
 
 use super::{cont_pages::ContPages, meta::PageMeta, Page};
-use crate::{boot::memory_region::MemoryRegionType, mm::PAGE_SIZE, sync::SpinLock};
+use crate::{
+    boot::memory_region::MemoryRegionType,
+    mm::{Paddr, PAGE_SIZE},
+    sync::SpinLock,
+};
 
 /// FrameAllocator with a counter for allocated memory
 pub(in crate::mm) struct CountingFrameAllocator {
@@ -58,26 +62,37 @@ impl CountingFrameAllocator {
 pub(in crate::mm) static PAGE_ALLOCATOR: Once<SpinLock<CountingFrameAllocator>> = Once::new();
 
 /// Allocate a single page.
-pub(crate) fn alloc_single<M: PageMeta>() -> Option<Page<M>> {
+///
+/// The metadata of the page is initialized with the given metadata.
+pub(crate) fn alloc_single<M: PageMeta>(metadata: M) -> Option<Page<M>> {
     PAGE_ALLOCATOR.get().unwrap().lock().alloc(1).map(|idx| {
         let paddr = idx * PAGE_SIZE;
-        Page::<M>::from_unused(paddr)
+        Page::from_unused(paddr, metadata)
     })
 }
 
 /// Allocate a contiguous range of pages of a given length in bytes.
 ///
+/// The caller must provide a closure to initialize metadata for all the pages.
+/// The closure receives the physical address of the page and returns the
+/// metadata, which is similar to [`core::array::from_fn`].
+///
 /// # Panics
 ///
 /// The function panics if the length is not base-page-aligned.
-pub(crate) fn alloc_contiguous<M: PageMeta>(len: usize) -> Option<ContPages<M>> {
+pub(crate) fn alloc_contiguous<M: PageMeta, F>(len: usize, metadata_fn: F) -> Option<ContPages<M>>
+where
+    F: FnMut(Paddr) -> M,
+{
     assert!(len % PAGE_SIZE == 0);
     PAGE_ALLOCATOR
         .get()
         .unwrap()
         .lock()
         .alloc(len / PAGE_SIZE)
-        .map(|start| ContPages::from_unused(start * PAGE_SIZE..start * PAGE_SIZE + len))
+        .map(|start| {
+            ContPages::from_unused(start * PAGE_SIZE..start * PAGE_SIZE + len, metadata_fn)
+        })
 }
 
 /// Allocate pages.
@@ -85,17 +100,24 @@ pub(crate) fn alloc_contiguous<M: PageMeta>(len: usize) -> Option<ContPages<M>>
 /// The allocated pages are not guarenteed to be contiguous.
 /// The total length of the allocated pages is `len`.
 ///
+/// The caller must provide a closure to initialize metadata for all the pages.
+/// The closure receives the physical address of the page and returns the
+/// metadata, which is similar to [`core::array::from_fn`].
+///
 /// # Panics
 ///
 /// The function panics if the length is not base-page-aligned.
-pub(crate) fn alloc<M: PageMeta>(len: usize) -> Option<Vec<Page<M>>> {
+pub(crate) fn alloc<M: PageMeta, F>(len: usize, mut metadata_fn: F) -> Option<Vec<Page<M>>>
+where
+    F: FnMut(Paddr) -> M,
+{
     assert!(len % PAGE_SIZE == 0);
     let nframes = len / PAGE_SIZE;
     let mut allocator = PAGE_ALLOCATOR.get().unwrap().lock();
     let mut vector = Vec::new();
     for _ in 0..nframes {
         let paddr = allocator.alloc(1)? * PAGE_SIZE;
-        let page = Page::<M>::from_unused(paddr);
+        let page = Page::<M>::from_unused(paddr, metadata_fn(paddr));
         vector.push(page);
     }
     Some(vector)

ostd/src/mm/page/cont_pages.rs

@@ -36,14 +36,21 @@ impl<M: PageMeta> Drop for ContPages<M> {
 impl<M: PageMeta> ContPages<M> {
     /// Create a new `ContPages` from unused pages.
     ///
+    /// The caller must provide a closure to initialize metadata for all the pages.
+    /// The closure receives the physical address of the page and returns the
+    /// metadata, which is similar to [`core::array::from_fn`].
+    ///
     /// # Panics
     ///
     /// The function panics if:
     ///  - the physical address is invalid or not aligned;
     ///  - any of the pages are already in use.
-    pub fn from_unused(range: Range<Paddr>) -> Self {
+    pub fn from_unused<F>(range: Range<Paddr>, mut metadata_fn: F) -> Self
+    where
+        F: FnMut(Paddr) -> M,
+    {
         for i in range.clone().step_by(PAGE_SIZE) {
-            let _ = ManuallyDrop::new(Page::<M>::from_unused(i));
+            let _ = ManuallyDrop::new(Page::<M>::from_unused(i, metadata_fn(i)));
         }
         Self {
             range,

ostd/src/mm/page/meta.rs

@@ -117,7 +117,7 @@ const_assert_eq!(size_of::<MetaSlot>(), 16);
 /// If a page type needs specific drop behavior, it should specify
 /// when implementing this trait. When we drop the last handle to
 /// this page, the `on_drop` method will be called.
-pub trait PageMeta: Default + Sync + private::Sealed + Sized {
+pub trait PageMeta: Sync + private::Sealed + Sized {
     const USAGE: PageUsage;
 
     fn on_drop(page: &mut Page<Self>);
@@ -184,7 +184,7 @@ pub struct PageTablePageMetaInner {
 
 /// The metadata of any kinds of page table pages.
 /// Make sure the the generic parameters don't effect the memory layout.
-#[derive(Debug, Default)]
+#[derive(Debug)]
 #[repr(C)]
 pub struct PageTablePageMeta<
     E: PageTableEntryTrait = PageTableEntry,
@@ -194,7 +194,23 @@ pub struct PageTablePageMeta<
 {
     pub lock: AtomicU8,
     pub inner: UnsafeCell<PageTablePageMetaInner>,
-    _phantom: core::marker::PhantomData<(E, C)>,
+    _phantom: PhantomData<(E, C)>,
+}
+
+impl<E: PageTableEntryTrait, C: PagingConstsTrait> PageTablePageMeta<E, C>
+where
+    [(); C::NR_LEVELS as usize]:,
+{
+    pub fn new_locked(level: PagingLevel) -> Self {
+        Self {
+            lock: AtomicU8::new(1),
+            inner: UnsafeCell::new(PageTablePageMetaInner {
+                level,
+                nr_children: 0,
+            }),
+            _phantom: PhantomData,
+        }
+    }
 }
 
 impl<E: PageTableEntryTrait, C: PagingConstsTrait> Sealed for PageTablePageMeta<E, C> where
@@ -275,7 +291,7 @@ pub(crate) fn init() -> Vec<Page<MetaPageMeta>> {
     // Now the metadata pages are mapped, we can initialize the metadata.
     meta_pages
         .into_iter()
-        .map(Page::<MetaPageMeta>::from_unused)
+        .map(|paddr| Page::<MetaPageMeta>::from_unused(paddr, MetaPageMeta::default()))
         .collect()
 }
 

ostd/src/mm/page/mod.rs

@@ -45,12 +45,14 @@ unsafe impl<M: PageMeta> Sync for Page<M> {}
 impl<M: PageMeta> Page<M> {
     /// Get a `Page` handle with a specific usage from a raw, unused page.
     ///
+    /// The caller should provide the initial metadata of the page.
+    ///
     /// # Panics
     ///
     /// The function panics if:
     ///  - the physical address is out of bound or not aligned;
     ///  - the page is already in use.
-    pub fn from_unused(paddr: Paddr) -> Self {
+    pub fn from_unused(paddr: Paddr, metadata: M) -> Self {
         assert!(paddr % PAGE_SIZE == 0);
         assert!(paddr < MAX_PADDR.load(Ordering::Relaxed) as Paddr);
         let vaddr = mapping::page_to_meta::<PagingConsts>(paddr);
@@ -72,7 +74,7 @@ impl<M: PageMeta> Page<M> {
         // SAFETY: The pointer points to the first byte of the `MetaSlot`
         // structure, and layout ensured enoungh space for `M`. The original
         // value does not represent any object that's needed to be dropped.
-        unsafe { (ptr as *mut M).write(M::default()) };
+        unsafe { (ptr as *mut M).write(metadata) };
 
         Self {
             ptr,

ostd/src/mm/page_table/node.rs

@@ -231,14 +231,8 @@ where
     /// set the lock bit for performance as it is exclusive and unlocking is an
     /// extra unnecessary expensive operation.
     pub(super) fn alloc(level: PagingLevel) -> Self {
-        let page = page::allocator::alloc_single::<PageTablePageMeta<E, C>>().unwrap();
+        let meta = PageTablePageMeta::new_locked(level);
-
+        let page = page::allocator::alloc_single::<PageTablePageMeta<E, C>>(meta).unwrap();
-        // The lock is initialized as held.
-        page.meta().lock.store(1, Ordering::Relaxed);
-
-        // SAFETY: here the page exclusively owned by the newly created handle.
-        let inner = unsafe { &mut *page.meta().inner.get() };
-        inner.level = level;
 
         // Zero out the page table node.
         let ptr = paddr_to_vaddr(page.paddr()) as *mut u8;

ostd/src/mm/page_table/test.rs

@@ -33,7 +33,7 @@ fn test_tracked_map_unmap() {
     let pt = PageTable::<UserMode>::empty();
 
     let from = PAGE_SIZE..PAGE_SIZE * 2;
-    let page = allocator::alloc_single::<FrameMeta>().unwrap();
+    let page = allocator::alloc_single(FrameMeta::default()).unwrap();
     let start_paddr = page.paddr();
     let prop = PageProperty::new(PageFlags::RW, CachePolicy::Writeback);
     unsafe { pt.cursor_mut(&from).unwrap().map(page.into(), prop) };
@@ -77,7 +77,7 @@ fn test_untracked_map_unmap() {
 fn test_user_copy_on_write() {
     let pt = PageTable::<UserMode>::empty();
     let from = PAGE_SIZE..PAGE_SIZE * 2;
-    let page = allocator::alloc_single::<FrameMeta>().unwrap();
+    let page = allocator::alloc_single(FrameMeta::default()).unwrap();
     let start_paddr = page.paddr();
     let prop = PageProperty::new(PageFlags::RW, CachePolicy::Writeback);
     unsafe { pt.cursor_mut(&from).unwrap().map(page.clone().into(), prop) };
@@ -133,7 +133,7 @@ fn test_base_protect_query() {
 
     let from_ppn = 1..1000;
     let from = PAGE_SIZE * from_ppn.start..PAGE_SIZE * from_ppn.end;
-    let to = allocator::alloc::<FrameMeta>(999 * PAGE_SIZE).unwrap();
+    let to = allocator::alloc(999 * PAGE_SIZE, |_| FrameMeta::default()).unwrap();
     let prop = PageProperty::new(PageFlags::RW, CachePolicy::Writeback);
     unsafe {
         let mut cursor = pt.cursor_mut(&from).unwrap();