[Patch v2] mm/huge_memory: use end_folio to terminate anonymous folio remapping

[Wei Yang <richard.weiyang@gmail.com>]

After splitting a large folio, it is necessary to remap the resulting
anonymous folios.

The current implementation determines the end of the remapping process
by counting the number of pages that have been processed.

Since the final folio in the sequence, end_folio, is already known and
tracked, this commit refactors the remapping loop to leverage end_folio
as the termination marker.

Signed-off-by: Wei Yang <richard.weiyang@gmail.com>
Cc: Zi Yan <ziy@nvidia.com>

---
v2: move folio assignment in loop
---
 mm/huge_memory.c | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 40cf59301c21..fe812d9c7807 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -3423,20 +3423,15 @@ bool unmap_huge_pmd_locked(struct vm_area_struct *vma, unsigned long addr,
 	return __discard_anon_folio_pmd_locked(vma, addr, pmdp, folio);
 }
 
-static void remap_page(struct folio *folio, unsigned long nr, int flags)
+static void remap_page(struct folio *folio, struct folio *end_folio, int flags)
 {
-	int i = 0;
-
 	/* If unmap_folio() uses try_to_migrate() on file, remove this check */
 	if (!folio_test_anon(folio))
 		return;
-	for (;;) {
+	do {
 		remove_migration_ptes(folio, folio, RMP_LOCKED | flags);
-		i += folio_nr_pages(folio);
-		if (i >= nr)
-			break;
 		folio = folio_next(folio);
-	}
+	} while (folio != end_folio);
 }
 
 static void lru_add_split_folio(struct folio *folio, struct folio *new_folio,
@@ -4066,7 +4061,7 @@ static int __folio_split(struct folio *folio, unsigned int new_order,
 	if (!ret && is_anon && !folio_is_device_private(folio))
 		remap_flags = RMP_USE_SHARED_ZEROPAGE;
 
-	remap_page(folio, 1 << old_order, remap_flags);
+	remap_page(folio, end_folio, remap_flags);
 
 	/*
 	 * Unlock all after-split folios except the one containing
-- 
2.34.1