From bb3c966f769ca3eb8ba30894c4b6ada64f89a4f4 Mon Sep 17 00:00:00 2001
From: Thamatip Chitpong <thamatip.chitpong@reactos.org>
Date: Thu, 21 Nov 2024 06:03:20 +0700
Subject: [PATCH] [NTOS:MM] section: Return number of pages for CC

- MmPurgeSegment
- MmMakeSegmentDirty
- MmFlushSegment
---
 ntoskrnl/mm/section.c | 61 +++++++++++++++++++++++++++++++------------
 1 file changed, 44 insertions(+), 17 deletions(-)

diff --git a/ntoskrnl/mm/section.c b/ntoskrnl/mm/section.c
index 99d5898781f9f..21ba7a555460e 100644
--- a/ntoskrnl/mm/section.c
+++ b/ntoskrnl/mm/section.c
@@ -4836,11 +4836,17 @@ NTAPI
 MmPurgeSegment(
     _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
     _In_opt_ PLARGE_INTEGER Offset,
-    _In_ ULONG Length)
+    _In_ ULONG Length,
+    _Out_opt_ PULONG PurgedDirtyPages)
 {
     LARGE_INTEGER PurgeStart, PurgeEnd;
     PMM_SECTION_SEGMENT Segment;
 
+    if (PurgedDirtyPages)
+    {
+        *PurgedDirtyPages = 0;
+    }
+
     PurgeStart.QuadPart = Offset ? Offset->QuadPart : 0LL;
     if (Length && Offset)
     {
@@ -4862,9 +4868,9 @@ MmPurgeSegment(
         /* We must calculate the length for ourselves */
         /* FIXME: All of this is suboptimal */
         ULONG ElemCount = RtlNumberGenericTableElements(&Segment->PageTable);
-        /* No page. Nothing to purge */
         if (!ElemCount)
         {
+            /* No page. Nothing to purge */
             MmUnlockSectionSegment(Segment);
             MmDereferenceSegment(Segment);
             return TRUE;
@@ -4877,6 +4883,7 @@ MmPurgeSegment(
     while (PurgeStart.QuadPart < PurgeEnd.QuadPart)
     {
         ULONG_PTR Entry = MmGetPageEntrySectionSegment(Segment, &PurgeStart);
+        BOOLEAN IsDirty;
 
         if (Entry == 0)
         {
@@ -4910,14 +4917,21 @@ MmPurgeSegment(
             return FALSE;
         }
 
+        IsDirty = IS_DIRTY_SSE(Entry);
+
         /* We can let this page go */
         MmSetPageEntrySectionSegment(Segment, &PurgeStart, 0);
         MmReleasePageMemoryConsumer(MC_USER, PFN_FROM_SSE(Entry));
 
+        /* Update number of purged dirty pages, if needed */
+        if (PurgedDirtyPages && IsDirty)
+        {
+            (*PurgedDirtyPages)++;
+        }
+
         PurgeStart.QuadPart += PAGE_SIZE;
     }
 
-    /* This page is currently in use. Bad luck */
     MmUnlockSectionSegment(Segment);
     MmDereferenceSegment(Segment);
     return TRUE;
@@ -4991,12 +5005,18 @@ NTAPI
 MmMakeSegmentDirty(
     _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
     _In_ LONGLONG Offset,
-    _In_ ULONG Length)
+    _In_ ULONG Length,
+    _Out_opt_ PULONG NumberOfPages)
 {
     PMM_SECTION_SEGMENT Segment;
     LARGE_INTEGER RangeStart, RangeEnd;
     NTSTATUS Status;
 
+    if (NumberOfPages)
+    {
+        *NumberOfPages = 0;
+    }
+
     RangeStart.QuadPart = Offset;
     Status = RtlLongLongAdd(RangeStart.QuadPart, Length, &RangeEnd.QuadPart);
     if (!NT_SUCCESS(Status))
@@ -5029,10 +5049,16 @@ MmMakeSegmentDirty(
         ASSERT(!IS_SWAP_FROM_SSE(Entry));
 
         /* If there is no page there, there is nothing to make dirty */
-        if (Entry != 0)
+        if (Entry != 0 && !IS_DIRTY_SSE(Entry))
         {
             /* Dirtify the entry */
             MmSetPageEntrySectionSegment(Segment, &RangeStart, DIRTY_SSE(Entry));
+
+            /* Update number of marked pages, if needed */
+            if (NumberOfPages)
+            {
+                (*NumberOfPages)++;
+            }
         }
 
         RangeStart.QuadPart += PAGE_SIZE;
@@ -5050,11 +5076,16 @@ MmFlushSegment(
     _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
     _In_opt_ PLARGE_INTEGER Offset,
     _In_ ULONG Length,
-    _Out_opt_ PIO_STATUS_BLOCK Iosb)
+    _Out_opt_ PULONG FlushedPages)
 {
     LARGE_INTEGER FlushStart, FlushEnd;
     NTSTATUS Status;
 
+    if (FlushedPages)
+    {
+        *FlushedPages = 0;
+    }
+
     if (Offset)
     {
         FlushStart = *Offset;
@@ -5063,14 +5094,11 @@ MmFlushSegment(
             return Status;
     }
 
-    if (Iosb)
-        Iosb->Information = 0;
-
     PMM_SECTION_SEGMENT Segment = MiGrabDataSection(SectionObjectPointer);
     if (!Segment)
     {
         /* Nothing to flush */
-        goto Quit;
+        return STATUS_SUCCESS;
     }
 
     ASSERT(*Segment->Flags & MM_DATAFILE_SEGMENT);
@@ -5088,7 +5116,7 @@ MmFlushSegment(
             /* No page. Nothing to flush */
             MmUnlockSectionSegment(Segment);
             MmDereferenceSegment(Segment);
-            goto Quit;
+            return STATUS_SUCCESS;
         }
 
         PCACHE_SECTION_PAGE_TABLE PageTable = RtlGetElementGenericTable(&Segment->PageTable, ElemCount - 1);
@@ -5106,8 +5134,11 @@ MmFlushSegment(
         {
             MmCheckDirtySegment(Segment, &FlushStart, FALSE, FALSE);
 
-            if (Iosb)
-                Iosb->Information += PAGE_SIZE;
+            /* Update number of flushed pages, if needed */
+            if (FlushedPages)
+            {
+                (*FlushedPages)++;
+            }
         }
 
         FlushStart.QuadPart += PAGE_SIZE;
@@ -5116,11 +5147,7 @@ MmFlushSegment(
     MmUnlockSectionSegment(Segment);
     MmDereferenceSegment(Segment);
 
-Quit:
     /* FIXME: Handle failures */
-    if (Iosb)
-        Iosb->Status = STATUS_SUCCESS;
-
     return STATUS_SUCCESS;
 }