From c7e09061caa02ee000ab33f1578bd3786e37fdd8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Gardou?= <jerome.gardou@reactos.org>
Date: Mon, 7 Jun 2021 14:12:21 +0200
Subject: [PATCH] [NTOS:MM] Get rid of MiQueryPageTableReferences

---
 ntoskrnl/mm/ARM3/miarm.h   | 29 -----------------------------
 ntoskrnl/mm/ARM3/virtual.c | 34 ++++++++++++++--------------------
 ntoskrnl/mm/i386/page.c    |  2 +-
 3 files changed, 15 insertions(+), 50 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/miarm.h b/ntoskrnl/mm/ARM3/miarm.h
index 2145b4c319e..721c66c54bb 100644
--- a/ntoskrnl/mm/ARM3/miarm.h
+++ b/ntoskrnl/mm/ARM3/miarm.h
@@ -2477,17 +2477,6 @@ MiDecrementPageTableReferences(IN PVOID Address)
     ASSERT(*RefCount < PTE_PER_PAGE);
     return *RefCount;
 }
-
-FORCEINLINE
-USHORT
-MiQueryPageTableReferences(IN PVOID Address)
-{
-    PUSHORT RefCount;
-
-    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
-
-    return *RefCount;
-}
 #else
 FORCEINLINE
 USHORT
@@ -2543,24 +2532,6 @@ MiDecrementPageTableReferences(IN PVOID Address)
 
     return Pfn->OriginalPte.u.Soft.UsedPageTableEntries;
 }
-
-FORCEINLINE
-USHORT
-MiQueryPageTableReferences(IN PVOID Address)
-{
-    PMMPDE PointerPde;
-    PMMPFN Pfn;
-
-    /* Make sure we're locked */
-    ASSERT((PsGetCurrentThread()->OwnsProcessWorkingSetExclusive) || (PsGetCurrentThread()->OwnsProcessWorkingSetShared));
-
-    PointerPde = MiAddressToPde(Address);
-    ASSERT(PointerPde->u.Hard.Valid);
-
-    /* This lies on the PFN */
-    Pfn = MiGetPfnEntry(PFN_FROM_PDE(PointerPde));
-    return Pfn->OriginalPte.u.Soft.UsedPageTableEntries;
-}
 #endif
 
 #ifdef __cplusplus
diff --git a/ntoskrnl/mm/ARM3/virtual.c b/ntoskrnl/mm/ARM3/virtual.c
index a94ba54afaa..27e2daa4241 100644
--- a/ntoskrnl/mm/ARM3/virtual.c
+++ b/ntoskrnl/mm/ARM3/virtual.c
@@ -659,12 +659,13 @@ MiDeleteVirtualAddresses(IN ULONG_PTR Va,
         PointerPte = MiAddressToPte(Va);
         do
         {
+            /* Making sure the PDE is still valid */
+            ASSERT(PointerPde->u.Hard.Valid == 1);
+
             /* Capture the PDE and make sure it exists */
             TempPte = *PointerPte;
             if (TempPte.u.Long)
             {
-                MiDecrementPageTableReferences((PVOID)Va);
-
                 /* Check if the PTE is actually mapped in */
                 if (MI_IS_MAPPED_PTE(&TempPte))
                 {
@@ -709,29 +710,23 @@ MiDeleteVirtualAddresses(IN ULONG_PTR Va,
                     /* The PTE was never mapped, just nuke it here */
                     MI_ERASE_PTE(PointerPte);
                 }
+
+                if (MiDecrementPageTableReferences((PVOID)Va) == 0)
+                {
+                    ASSERT(PointerPde->u.Long != 0);
+                    /* Delete the PDE proper */
+                    MiDeletePde(PointerPde, CurrentProcess);
+                    /* Jump */
+                    Va = (ULONG_PTR)MiPdeToAddress(PointerPde + 1);
+                    break;
+                }
             }
 
             /* Update the address and PTE for it */
             Va += PAGE_SIZE;
             PointerPte++;
             PrototypePte++;
-
-            /* Making sure the PDE is still valid */
-            ASSERT(PointerPde->u.Hard.Valid == 1);
-        }
-        while ((Va & (PDE_MAPPED_VA - 1)) && (Va <= EndingAddress));
-
-        /* The PDE should still be valid at this point */
-        ASSERT(PointerPde->u.Hard.Valid == 1);
-
-        /* Check remaining PTE count (go back 1 page due to above loop) */
-        if (MiQueryPageTableReferences((PVOID)(Va - PAGE_SIZE)) == 0)
-        {
-            ASSERT(PointerPde->u.Long != 0);
-
-            /* Delete the PDE proper */
-            MiDeletePde(PointerPde, CurrentProcess);
-        }
+        } while ((Va & (PDE_MAPPED_VA - 1)) && (Va <= EndingAddress));
 
         /* Release the lock */
         MiReleasePfnLock(OldIrql);
@@ -739,7 +734,6 @@ MiDeleteVirtualAddresses(IN ULONG_PTR Va,
         if (Va > EndingAddress) return;
 
         /* Otherwise, we exited because we hit a new PDE boundary, so start over */
-        PointerPde = MiAddressToPde(Va);
         AddressGap = FALSE;
     }
 }
diff --git a/ntoskrnl/mm/i386/page.c b/ntoskrnl/mm/i386/page.c
index 2c402c51478..7210d769df8 100644
--- a/ntoskrnl/mm/i386/page.c
+++ b/ntoskrnl/mm/i386/page.c
@@ -119,7 +119,7 @@ BOOLEAN
 MiIsPageTablePresent(PVOID Address)
 {
 #if _MI_PAGING_LEVELS == 2
-    return MiQueryPageTableReferences(Address) != 0;
+    return MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)] != 0;
 #else
     PMMPDE PointerPde;
     PMMPPE PointerPpe;