Wrap mutable data into UnsafeCell

ostd/src/mm/page/meta.rs

@@ -37,6 +37,7 @@ pub mod mapping {
 
 use alloc::vec::Vec;
 use core::{
+    cell::UnsafeCell,
     marker::PhantomData,
     mem::{size_of, ManuallyDrop},
     panic,
@@ -174,6 +175,13 @@ pub struct FrameMeta {
 
 impl Sealed for FrameMeta {}
 
+#[derive(Debug, Default)]
+#[repr(C)]
+pub struct PageTablePageMetaInner {
+    pub level: PagingLevel,
+    pub nr_children: u16,
+}
+
 /// The metadata of any kinds of page table pages.
 /// Make sure the the generic parameters don't effect the memory layout.
 #[derive(Debug, Default)]
@@ -185,8 +193,7 @@ pub struct PageTablePageMeta<
     [(); C::NR_LEVELS as usize]:,
 {
     pub lock: AtomicU8,
-    pub level: PagingLevel,
+    pub inner: UnsafeCell<PageTablePageMetaInner>,
-    pub nr_children: u16,
     _phantom: core::marker::PhantomData<(E, C)>,
 }
 
@@ -195,6 +202,16 @@ impl<E: PageTableEntryTrait, C: PagingConstsTrait> Sealed for PageTablePageMeta<
 {
 }
 
+unsafe impl<E: PageTableEntryTrait, C: PagingConstsTrait> Send for PageTablePageMeta<E, C> where
+    [(); C::NR_LEVELS as usize]:
+{
+}
+
+unsafe impl<E: PageTableEntryTrait, C: PagingConstsTrait> Sync for PageTablePageMeta<E, C> where
+    [(); C::NR_LEVELS as usize]:
+{
+}
+
 #[derive(Debug, Default)]
 #[repr(C)]
 pub struct MetaPageMeta {}

ostd/src/mm/page/mod.rs

@@ -142,15 +142,6 @@ impl<M: PageMeta> Page<M> {
         unsafe { &*(self.ptr as *const M) }
     }
 
-    /// Get the mutable metadata of this page.
-    ///
-    /// # Safety
-    ///
-    /// The caller should be sure that the page is exclusively owned.
-    pub(in crate::mm) unsafe fn meta_mut(&mut self) -> &mut M {
-        unsafe { &mut *(self.ptr as *mut M) }
-    }
-
     fn get_ref_count(&self) -> &AtomicU32 {
         unsafe { &(*self.ptr).ref_count }
     }

ostd/src/mm/page_table/node.rs

@@ -34,7 +34,7 @@ use crate::{
         paddr_to_vaddr,
         page::{
             self,
-            meta::{PageMeta, PageTablePageMeta, PageUsage},
+            meta::{PageMeta, PageTablePageMeta, PageTablePageMetaInner, PageUsage},
             DynPage, Page,
         },
         page_prop::PageProperty,
@@ -208,13 +208,14 @@ 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 mut page = page::allocator::alloc_single::<PageTablePageMeta<E, C>>().unwrap();
+        let page = page::allocator::alloc_single::<PageTablePageMeta<E, C>>().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.
-        unsafe { page.meta_mut().level = level };
+        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;
@@ -227,7 +228,7 @@ where
     }
 
     pub fn level(&self) -> PagingLevel {
-        self.page.meta().level
+        self.meta().level
     }
 
     /// Converts the handle into a raw handle to be stored in a PTE or CPU.
@@ -478,20 +479,31 @@ where
 
             // Update the child count.
             if pte.is_none() {
-                // SAFETY: Here we have an exclusive access to the page.
+                self.meta_mut().nr_children -= 1;
-                unsafe { self.page.meta_mut().nr_children -= 1 };
             }
         } else if let Some(e) = pte {
             // SAFETY: This is safe as described in the above branch.
             unsafe { (self.as_ptr() as *mut E).add(idx).write(e) };
-            // SAFETY: Here we have an exclusive access to the page.
+
-            unsafe { self.page.meta_mut().nr_children += 1 };
+            self.meta_mut().nr_children += 1;
         }
     }
 
     fn as_ptr(&self) -> *const E {
         paddr_to_vaddr(self.start_paddr()) as *const E
     }
+
+    fn meta(&self) -> &PageTablePageMetaInner {
+        // SAFETY: We have exclusively locked the page, so we can derive an immutable reference
+        // from an immutable reference to the lock guard.
+        unsafe { &*self.page.meta().inner.get() }
+    }
+
+    fn meta_mut(&mut self) -> &mut PageTablePageMetaInner {
+        // SAFETY: We have exclusively locked the page, so we can derive a mutable reference from a
+        // mutable reference to the lock guard.
+        unsafe { &mut *self.page.meta().inner.get() }
+    }
 }
 
 impl<E: PageTableEntryTrait, C: PagingConstsTrait> Drop for PageTableNode<E, C>
@@ -512,7 +524,9 @@ where
 
     fn on_drop(page: &mut Page<Self>) {
         let paddr = page.paddr();
-        let level = page.meta().level;
+
+        let inner = unsafe { &mut *page.meta().inner.get() };
+        let level = inner.level;
 
         // Drop the children.
         for i in 0..nr_subpage_per_huge::<C>() {