Prefer ManuallyDrop than mem::forget

ostd/src/mm/page/mod.rs

@@ -116,6 +116,17 @@ impl<M: PageMeta> Page<M> {
         }
     }
 
+    /// Increase the reference count of the page by one.
+    ///
+    /// # Safety
+    ///
+    /// The physical address must represent a valid page and the caller must already hold one
+    /// reference count.
+    pub(in crate::mm) unsafe fn inc_ref(paddr: Paddr) {
+        let page = unsafe { ManuallyDrop::new(Self::from_raw(paddr)) };
+        let _page = page.clone();
+    }
+
     /// Get the physical address.
     pub fn paddr(&self) -> Paddr {
         mapping::meta_to_page::<PagingConsts>(self.ptr as Vaddr)
@@ -211,6 +222,17 @@ impl DynPage {
         Self { ptr }
     }
 
+    /// Increase the reference count of the page by one.
+    ///
+    /// # Safety
+    ///
+    /// The physical address must represent a valid page and the caller must already hold one
+    /// reference count.
+    pub(in crate::mm) unsafe fn inc_ref(paddr: Paddr) {
+        let page = unsafe { ManuallyDrop::new(Self::from_raw(paddr)) };
+        let _page = page.clone();
+    }
+
     /// Get the physical address of the start of the page
     pub fn paddr(&self) -> Paddr {
         mapping::meta_to_page::<PagingConsts>(self.ptr as Vaddr)

ostd/src/mm/page_table/node.rs

@@ -69,10 +69,13 @@ where
 
     /// Converts a raw handle to an accessible handle by pertaining the lock.
     pub(super) fn lock(self) -> PageTableNode<E, C> {
+        // Prevent dropping the handle.
+        let this = ManuallyDrop::new(self);
+
         // SAFETY: The physical address in the raw handle is valid and we are
         // transferring the ownership to a new handle. No increment of the reference
         // count is needed.
-        let page = unsafe { Page::<PageTablePageMeta<E, C>>::from_raw(self.paddr()) };
+        let page = unsafe { Page::<PageTablePageMeta<E, C>>::from_raw(this.paddr()) };
 
         // Acquire the lock.
         while page
@@ -84,9 +87,6 @@ where
             core::hint::spin_loop();
         }
 
-        // Prevent dropping the handle.
-        let _ = ManuallyDrop::new(self);
-
         PageTableNode::<E, C> { page }
     }
 
@@ -149,11 +149,11 @@ where
     }
 
     fn inc_ref(&self) {
-        // SAFETY: The physical address in the raw handle is valid and we are
+        // SAFETY: We have a reference count to the page and can safely increase the reference
-        // incrementing the reference count by cloning and forgetting.
+        // count by one more.
-        let page = unsafe { Page::<PageTablePageMeta<E, C>>::from_raw(self.paddr()) };
+        unsafe {
-        core::mem::forget(page.clone());
+            Page::<PageTablePageMeta<E, C>>::inc_ref(self.paddr());
-        core::mem::forget(page);
+        }
     }
 }
 
@@ -264,24 +264,22 @@ where
         } else {
             let paddr = pte.paddr();
             if !pte.is_last(self.level()) {
-                // SAFETY: The physical address is recorded in a valid PTE
+                // SAFETY: We have a reference count to the page and can safely increase the reference
-                // which would be casted from a handle. We are incrementing
+                // count by one more.
-                // the reference count so we restore, clone, and forget both.
+                unsafe {
-                let node = unsafe { Page::<PageTablePageMeta<E, C>>::from_raw(paddr) };
+                    Page::<PageTablePageMeta<E, C>>::inc_ref(paddr);
-                let inc_ref = node.clone();
+                }
-                core::mem::forget(node);
-                core::mem::forget(inc_ref);
                 Child::PageTable(RawPageTableNode {
                     raw: paddr,
                     _phantom: PhantomData,
                 })
             } else if in_tracked_range {
-                // SAFETY: The physical address is recorded in a valid PTE
+                // SAFETY: We have a reference count to the page and can safely increase the reference
-                // which would be casted from a handle. We are incrementing
+                // count by one more.
-                // the reference count so we restore and forget a cloned one.
+                unsafe {
-                let page = unsafe { DynPage::from_raw(paddr) };
+                    DynPage::inc_ref(paddr);
-                core::mem::forget(page.clone());
+                }
-                Child::Page(page)
+                Child::Page(unsafe { DynPage::from_raw(paddr) })
             } else {
                 Child::Untracked(paddr)
             }
@@ -358,10 +356,11 @@ where
         // They should be ensured by the cursor.
         debug_assert!(idx < nr_subpage_per_huge::<C>());
 
+        // The ownership is transferred to a raw PTE. Don't drop the handle.
+        let pt = ManuallyDrop::new(pt);
+
         let pte = Some(E::new_pt(pt.paddr()));
         self.overwrite_pte(idx, pte, in_tracked_range);
-        // The ownership is transferred to a raw PTE. Don't drop the handle.
-        let _ = ManuallyDrop::new(pt);
     }
 
     /// Map a page at a given index.