From c05aeb95c3610dbecce5bc40e9c561c4e5572710 Mon Sep 17 00:00:00 2001
From: Laurence Rowe <l@lrowe.co.uk>
Date: Tue, 4 Nov 2025 01:10:12 -0800
Subject: [PATCH] Fix get_accessed_pages to include 4KB pages.

---
 lib/tinykvm/amd64/paging.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/tinykvm/amd64/paging.cpp b/lib/tinykvm/amd64/paging.cpp
index 8c7ccaa..ccb609c 100644
--- a/lib/tinykvm/amd64/paging.cpp
+++ b/lib/tinykvm/amd64/paging.cpp
@@ -552,7 +552,9 @@ std::vector<std::pair<uint64_t, uint64_t>> get_accessed_pages(const vMemory& mem
 	std::vector<std::pair<uint64_t, uint64_t>> accessed_pages;
 	foreach_page(memory,
 	[&accessed_pages] (uint64_t addr, uint64_t& entry, size_t size) {
-		if ((entry & (PDE64_PS | PDE64_ACCESSED | PDE64_PRESENT)) == (PDE64_PS | PDE64_ACCESSED | PDE64_PRESENT)) {
+		// Leaf pages are either huge pages with the PS bit set or of PAGE_SIZE.
+		if ((entry & (PDE64_ACCESSED | PDE64_PRESENT)) == (PDE64_ACCESSED | PDE64_PRESENT) &&
+				((entry & PDE64_PS) || (size == PAGE_SIZE))) {
 			accessed_pages.push_back({addr & PDE64_ADDR_MASK, size});
 		}
 	}, false);