From 00d9b01d5c356fbc5f6b72b9a69a7571fdf0ab1f Mon Sep 17 00:00:00 2001
From: Chen Chengjun <chenchengjun.ccj@antgroup.com>
Date: Thu, 29 Aug 2024 11:52:33 +0800
Subject: [PATCH] Add a fast path to page table fork

---
 ostd/src/mm/page_table/node.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/ostd/src/mm/page_table/node.rs b/ostd/src/mm/page_table/node.rs
index b6da9b8df..37e54ca74 100644
--- a/ostd/src/mm/page_table/node.rs
+++ b/ostd/src/mm/page_table/node.rs
@@ -355,17 +355,22 @@ where
         debug_assert!(deep.end <= shallow.start || deep.start >= shallow.end);
 
         let mut new_pt = Self::alloc(self.level());
-
+        let mut copied_child_count = self.nr_children();
         for i in deep {
+            if copied_child_count == 0 {
+                return new_pt;
+            }
             match self.child(i, true) {
                 Child::PageTable(pt) => {
                     let guard = pt.clone_shallow().lock();
                     let new_child = guard.make_copy(0..nr_subpage_per_huge::<C>(), 0..0);
                     new_pt.set_child_pt(i, new_child.into_raw(), true);
+                    copied_child_count -= 1;
                 }
                 Child::Page(page) => {
                     let prop = self.read_pte_prop(i);
                     new_pt.set_child_page(i, page.clone(), prop);
+                    copied_child_count -= 1;
                 }
                 Child::None => {}
                 Child::Untracked(_) => {
@@ -375,10 +380,14 @@ where
         }
 
         for i in shallow {
+            if copied_child_count == 0 {
+                return new_pt;
+            }
             debug_assert_eq!(self.level(), C::NR_LEVELS);
             match self.child(i, /*meaningless*/ true) {
                 Child::PageTable(pt) => {
                     new_pt.set_child_pt(i, pt.clone_shallow(), /*meaningless*/ true);
+                    copied_child_count -= 1;
                 }
                 Child::None => {}
                 Child::Page(_) | Child::Untracked(_) => {