Index: ntoskrnl/mm/section.c
===================================================================
--- ntoskrnl/mm/section.c	(revision 63990)
+++ ntoskrnl/mm/section.c	(working copy)
@@ -177,7 +177,88 @@
 
 /* FUNCTIONS *****************************************************************/
 
+NTSTATUS NTAPI
+MiCreateVirtualMappingAndUpdatePdeCount(PEPROCESS Process,
+										PVOID Address,
+										ULONG Protect,
+										PPFN_NUMBER Pages,
+										ULONG Count,
+										BOOLEAN Unsafe)
+{
+    NTSTATUS Status = (Unsafe ? &MmCreateVirtualMappingUnsafe : &MmCreateVirtualMapping)
+        (Process, Address, Protect, Pages, Count);
+    /* Update pde ref counts -- MmCreateVirtualMapping requires the old PTE to be zero. */
+#if (_MI_PAGING_LEVELS == 2)
+    /* Reference Page Directory Entry */
+    while (Count--)
+    {
+        if (Address < MmSystemRangeStart)
+	{
+            MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)]++;
+            ASSERT(MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)] <= PTE_COUNT);
+        }
+	Address = (PVOID)((PCHAR)Address + PAGE_SIZE);
+    }
+#endif
+    return Status;
+}
 
+VOID NTAPI
+MiDeleteVirtualMappingAndUpdatePdeCount(PEPROCESS Process,
+										PVOID Address,
+										BOOLEAN FreePage,
+										PBOOLEAN WasDirty,
+										PPFN_NUMBER OldPage)
+{
+    MiDeleteVirtualMappingAndUpdatePdeCount
+        (Process, Address, FreePage, WasDirty, OldPage);
+    /* By definition, the page entry will not be used after this */
+#if (_MI_PAGING_LEVELS == 2)
+    if (Address < MmSystemRangeStart)
+    {
+        MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)]--;
+        ASSERT(MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)] <= PTE_COUNT);
+    }
+#endif
+}
+
+NTSTATUS NTAPI
+MiCreatePageFileMappingAndUpdatePdeCount(PEPROCESS Process,
+										PVOID Address,
+										SWAPENTRY SwapEntry)
+{
+    NTSTATUS Status = MmCreatePageFileMapping(Process, Address, SwapEntry);
+    /* Update pde ref counts -- MmCreatePageFileMapping requires the old PTE to be zero. */
+#if (_MI_PAGING_LEVELS == 2)
+    if (Address < MmSystemRangeStart)
+    {
+        MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)]++;
+        ASSERT(MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)] <= PTE_COUNT);
+    }
+#endif
+    return Status;
+}
+
+VOID NTAPI
+MiDeletePageFileMappingAndUpdatePdeCount(PEPROCESS Process,
+										PVOID Address,
+										SWAPENTRY *SwapEntry)
+{
+    MmDeletePageFileMapping(Process, Address, SwapEntry);
+#if (_MI_PAGING_LEVELS == 2)
+    if (*SwapEntry && Address < MmSystemRangeStart)
+    {
+        MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)]--;
+        ASSERT(MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)] <= PTE_COUNT);
+    }
+#endif
+}
+
+//#define MmCreateVirtualMapping __Please_use_MiCreateVirtualMappingAndUpdatePdeCount_instead
+//#define MmDeleteVirtualMapping __Please_use_MiDeleteVirtualMappingAndUpdatePdeCount_instead
+//#define MmCreatePageFileMapping __Please_use_MiCreatePageFileMappingAndUpdatePdeCount_instead
+//#define MmDeletePageFileMapping __Please_use_MiDeletePageFileMappingAndUpdatePdeCount_instead
+
 /*
  References:
   [1] Microsoft Corporation, "Microsoft Portable Executable and Common Object
@@ -1396,7 +1477,7 @@
       }
 
       MmUnlockSectionSegment(Segment);
-      MmDeletePageFileMapping(Process, Address, &SwapEntry);
+      MiDeletePageFileMappingAndUpdatePdeCount(Process, Address, &SwapEntry);
       MmCreatePageFileMapping(Process, Address, MM_WAIT_ENTRY);
 
       MmUnlockAddressSpace(AddressSpace);
@@ -1417,14 +1498,15 @@
       }
       MmLockAddressSpace(AddressSpace);
       MmDeletePageFileMapping(Process, PAddress, &DummyEntry);
-      Status = MmCreateVirtualMapping(Process,
+      Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                       PAddress,
                                       Region->Protect,
                                       &Page,
-                                      1);
+                                      1,
+									  FALSE);
       if (!NT_SUCCESS(Status))
       {
-         DPRINT("MmCreateVirtualMapping failed, not out of memory\n");
+         DPRINT("MiCreateVirtualMappingAndUpdatePdeCount failed, not out of memory\n");
           KeBugCheck(MEMORY_MANAGEMENT);
          return(Status);
       }
@@ -1456,14 +1538,15 @@
        * Just map the desired physical page
        */
       Page = (PFN_NUMBER)(Offset.QuadPart >> PAGE_SHIFT);
-      Status = MmCreateVirtualMappingUnsafe(Process,
+      Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                             PAddress,
                                             Region->Protect,
                                             &Page,
-                                            1);
+                                            1,
+											TRUE);
       if (!NT_SUCCESS(Status))
       {
-         DPRINT("MmCreateVirtualMappingUnsafe failed, not out of memory\n");
+         DPRINT("MiCreateVirtualMappingAndUpdatePdeCount failed, not out of memory\n");
           KeBugCheck(MEMORY_MANAGEMENT);
          return(Status);
       }
@@ -1547,11 +1630,12 @@
       MmDeletePageFileMapping(Process, PAddress, &FakeSwapEntry);
       DPRINT("CreateVirtualMapping Page %x Process %p PAddress %p Attributes %x\n",
               Page, Process, PAddress, Attributes);
-      Status = MmCreateVirtualMapping(Process,
+      Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                       PAddress,
                                       Attributes,
                                       &Page,
-                                      1);
+                                      1,
+									  FALSE);
       if (!NT_SUCCESS(Status))
       {
          DPRINT1("Unable to create virtual mapping\n");
@@ -1620,11 +1704,12 @@
        * Save the swap entry.
        */
       MmSetSavedSwapEntryPage(Page, SwapEntry);
-      Status = MmCreateVirtualMapping(Process,
+      Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                       PAddress,
                                       Region->Protect,
                                       &Page,
-                                      1);
+                                      1,
+									  FALSE);
       if (!NT_SUCCESS(Status))
       {
          DPRINT1("Unable to create virtual mapping\n");
@@ -1647,11 +1732,12 @@
       MmSharePageEntrySectionSegment(Segment, &Offset);
       MmUnlockSectionSegment(Segment);
 
-      Status = MmCreateVirtualMapping(Process,
+      Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                       PAddress,
                                       Attributes,
                                       &Page,
-                                      1);
+                                      1,
+									  FALSE);
       if (!NT_SUCCESS(Status))
       {
          DPRINT1("Unable to create virtual mapping\n");
@@ -1761,7 +1847,7 @@
    }
 
    MmDeleteRmap(OldPage, Process, PAddress);
-   MmDeleteVirtualMapping(Process, PAddress, FALSE, NULL, NULL);
+   MiDeleteVirtualMappingAndUpdatePdeCount(Process, PAddress, FALSE, NULL, NULL);
    MmCreatePageFileMapping(Process, PAddress, MM_WAIT_ENTRY);
 
    /*
@@ -1793,14 +1879,15 @@
     * Set the PTE to point to the new page
     */
    MmDeletePageFileMapping(Process, PAddress, &SwapEntry);
-   Status = MmCreateVirtualMapping(Process,
+   Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                    PAddress,
                                    Region->Protect,
                                    &NewPage,
-                                   1);
+                                   1,
+								   FALSE);
     if (!NT_SUCCESS(Status))
     {
-        DPRINT1("MmCreateVirtualMapping failed, unable to create virtual mapping, not out of memory\n");
+        DPRINT1("MiCreateVirtualMappingAndUpdatePdeCount failed, unable to create virtual mapping, not out of memory\n");
         KeBugCheck(MEMORY_MANAGEMENT);
         return(Status);
     }
@@ -1832,7 +1919,7 @@
       MmLockAddressSpace(&Process->Vm);
    }
 
-   MmDeleteVirtualMapping(Process,
+   MiDeleteVirtualMappingAndUpdatePdeCount(Process,
                           Address,
                           FALSE,
                           &WasDirty,
@@ -2104,7 +2191,7 @@
       DPRINT("Not dirty and private and not swapped (%p:%p)\n", Process, Address);
       MmSetSavedSwapEntryPage(Page, 0);
       MmLockAddressSpace(AddressSpace);
-      Status = MmCreatePageFileMapping(Process,
+      Status = MiCreatePageFileMappingAndUpdatePdeCount(Process,
                                        Address,
                                        SwapEntry);
       MmUnlockAddressSpace(AddressSpace);
@@ -2133,11 +2220,12 @@
           */
          if (Context.Private)
          {
-            Status = MmCreateVirtualMapping(Process,
+            Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                             Address,
                                             MemoryArea->Protect,
                                             &Page,
-                                            1);
+                                            1,
+											FALSE);
             MmSetDirtyPage(Process, Address);
             MmInsertRmap(Page,
                          Process,
@@ -2151,11 +2239,12 @@
              * set it back into the section segment entry so we don't loose
              * our copy. Otherwise it will be handled by the cache manager.
              */
-            Status = MmCreateVirtualMapping(Process,
+            Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                             Address,
                                             MemoryArea->Protect,
                                             &Page,
-                                            1);
+                                            1,
+											FALSE);
             MmSetDirtyPage(Process, Address);
             MmInsertRmap(Page,
                          Process,
@@ -2189,11 +2278,12 @@
       MmLockAddressSpace(AddressSpace);
       if (Context.Private)
       {
-         Status = MmCreateVirtualMapping(Process,
+         Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                          Address,
                                          MemoryArea->Protect,
                                          &Page,
-                                         1);
+                                         1,
+										 FALSE);
          MmSetDirtyPage(Process, Address);
          MmInsertRmap(Page,
                       Process,
@@ -2201,11 +2291,12 @@
       }
       else
       {
-         Status = MmCreateVirtualMapping(Process,
+         Status = MiCreateVirtualMappingAndUpdatePdeCount(Process,
                                          Address,
                                          MemoryArea->Protect,
                                          &Page,
-                                         1);
+                                         1,
+										 FALSE);
          MmSetDirtyPage(Process, Address);
          MmInsertRmap(Page,
                       Process,
@@ -2239,7 +2330,7 @@
    {
       MmLockAddressSpace(AddressSpace);
       MmLockSectionSegment(Context.Segment);
-      Status = MmCreatePageFileMapping(Process,
+      Status = MiCreatePageFileMappingAndUpdatePdeCount(Process,
                                        Address,
                                        SwapEntry);
       /* We had placed a wait entry upon entry ... replace it before leaving */