drgn.helpers.linux.mm: add compound page helpers

I had these helpers lying around from a couple of bugs related to
compound pages that I debugged.

Signed-off-by: Omar Sandoval <osandov@osandov.com>

drgn/helpers/linux/mm.py

@@ -20,12 +20,19 @@
 __all__ = (
     "PFN_PHYS",
     "PHYS_PFN",
+    "PageCompound",
+    "PageHead",
+    "PageTail",
     "access_process_vm",
     "access_remote_vm",
     "cmdline",
+    "compound_head",
+    "compound_nr",
+    "compound_order",
     "decode_page_flags",
     "environ",
     "for_each_page",
+    "page_size",
     "page_to_pfn",
     "page_to_phys",
     "page_to_virt",
@@ -558,6 +565,113 @@ def PageYoung(page: Object) -> bool:
     return bool(page.flags & (1 << flag))
 
 
+# End generated by scripts/generate_page_flag_getters.py.
+
+
+def PageCompound(page: Object) -> bool:
+    """
+    Return whether a page is part of a `compound page
+    <https://lwn.net/Articles/619514/>`_.
+
+    :param page: ``struct page *``
+    """
+    page = page.read_()
+    return bool(
+        (page.flags & (1 << page.prog_["PG_head"])) or (page.compound_head.read_() & 1)
+    )
+
+
+# HugeTLB Vmemmap Optimization (HVO) creates "fake" head pages that are
+# actually tail pages. See Linux kernel commit e7d324850bfc ("mm: hugetlb: free
+# the 2nd vmemmap page associated with each HugeTLB page") (in v5.18) and
+# https://www.kernel.org/doc/html/latest/mm/vmemmap_dedup.html.
+def _page_is_fake_head(page: Object) -> bool:
+    head = page[1].compound_head.value_()
+    return bool(head & 1) and (head - 1) != page.value_()
+
+
+def PageHead(page: Object) -> bool:
+    """
+    Return whether a page is a head page in a `compound page`_.
+
+    :param page: ``struct page *``
+    """
+    page = page.read_()
+    return bool(
+        page.flags & (1 << page.prog_["PG_head"]) and not _page_is_fake_head(page)
+    )
+
+
+def PageTail(page: Object) -> bool:
+    """
+    Return whether a page is a tail page in a `compound page`_.
+
+    :param page: ``struct page *``
+    """
+    page = page.read_()
+    if page.compound_head.value_() & 1:
+        return True
+    if page.flags & (1 << page.prog_["PG_head"]):
+        return _page_is_fake_head(page)
+    return False
+
+
+def compound_head(page: Object) -> Object:
+    """
+    Get the head page associated with a page.
+
+    If *page* is a tail page, this returns the head page of the `compound
+    page`_ it belongs to. Otherwise, it returns *page*.
+
+    :param page: ``struct page *``
+    :return: ``struct page *``
+    """
+    page = page.read_()
+    head = page.compound_head.read_()
+    if head & 1:
+        return cast(page.type_, head - 1)
+    # Handle fake head pages (see _page_is_fake_head()).
+    if page.flags & (1 << page.prog_["PG_head"]):
+        head = page[1].compound_head.read_()
+        if head & 1:
+            return cast(page.type_, head - 1)
+    return page
+
+
+def compound_order(page: Object) -> Object:
+    """
+    Return the allocation order of a potentially `compound page`_.
+
+    :param page: ``struct page *``
+    :return: ``unsigned int``
+    """
+    if not PageHead(page):
+        return Object(page.prog_, "unsigned int", 0)
+    return cast("unsigned int", page[1].compound_order)
+
+
+def compound_nr(page: Object) -> Object:
+    """
+    Return the number of pages in a potentially `compound page`_.
+
+    :param page: ``struct page *``
+    :return: ``unsigned long``
+    """
+    if not PageHead(page):
+        return Object(page.prog_, "unsigned long", 1)
+    return Object(page.prog_, "unsigned long", 1) << page[1].compound_order
+
+
+def page_size(page: Object) -> Object:
+    """
+    Return the number of bytes in a potentially `compound page`_.
+
+    :param page: ``struct page *``
+    :return: ``unsigned long``
+    """
+    return page.prog_["PAGE_SIZE"] << compound_order(page)
+
+
 def decode_page_flags(page: Object) -> str:
     """
     Get a human-readable representation of the flags set on a page.

scripts/generate_page_flag_getters.py

@@ -45,3 +45,5 @@ def Page{uname}(page: Object) -> bool:
     return bool(page.flags & (1 << flag))
 '''
         )
+    print()
+    print("# End generated by scripts/generate_page_flag_getters.py.")

tests/linux_kernel/helpers/test_mm.py

@@ -14,13 +14,20 @@
 from drgn.helpers.linux.mm import (
     PFN_PHYS,
     PHYS_PFN,
+    PageCompound,
+    PageHead,
     PageSwapBacked,
+    PageTail,
     PageWriteback,
     access_process_vm,
     access_remote_vm,
     cmdline,
+    compound_head,
+    compound_nr,
+    compound_order,
     decode_page_flags,
     environ,
+    page_size,
     page_to_pfn,
     page_to_phys,
     page_to_virt,
@@ -80,6 +87,55 @@ def test_page_flag_getters(self):
             self.assertTrue(PageSwapBacked(page))
             self.assertFalse(PageWriteback(page))
 
+    @skip_unless_have_test_kmod
+    def test_PageCompound(self):
+        self.assertFalse(PageCompound(self.prog["drgn_test_page"]))
+        self.assertTrue(PageCompound(self.prog["drgn_test_compound_page"]))
+        self.assertTrue(PageCompound(self.prog["drgn_test_compound_page"] + 1))
+
+    @skip_unless_have_test_kmod
+    def test_PageHead(self):
+        self.assertFalse(PageHead(self.prog["drgn_test_page"]))
+        self.assertTrue(PageHead(self.prog["drgn_test_compound_page"]))
+        self.assertFalse(PageHead(self.prog["drgn_test_compound_page"] + 1))
+
+    @skip_unless_have_test_kmod
+    def test_PageTail(self):
+        self.assertFalse(PageTail(self.prog["drgn_test_page"]))
+        self.assertFalse(PageTail(self.prog["drgn_test_compound_page"]))
+        self.assertTrue(PageTail(self.prog["drgn_test_compound_page"] + 1))
+
+    @skip_unless_have_test_kmod
+    def test_compound_head(self):
+        self.assertEqual(
+            compound_head(self.prog["drgn_test_page"]), self.prog["drgn_test_page"]
+        )
+        self.assertEqual(
+            compound_head(self.prog["drgn_test_compound_page"]),
+            self.prog["drgn_test_compound_page"],
+        )
+        self.assertEqual(
+            compound_head(self.prog["drgn_test_compound_page"] + 1),
+            self.prog["drgn_test_compound_page"],
+        )
+
+    @skip_unless_have_test_kmod
+    def test_compound_order(self):
+        self.assertEqual(compound_order(self.prog["drgn_test_page"]), 0)
+        self.assertEqual(compound_order(self.prog["drgn_test_compound_page"]), 1)
+
+    @skip_unless_have_test_kmod
+    def test_compound_nr(self):
+        self.assertEqual(compound_nr(self.prog["drgn_test_page"]), 1)
+        self.assertEqual(compound_nr(self.prog["drgn_test_compound_page"]), 2)
+
+    @skip_unless_have_test_kmod
+    def test_page_size(self):
+        self.assertEqual(page_size(self.prog["drgn_test_page"]), self.prog["PAGE_SIZE"])
+        self.assertEqual(
+            page_size(self.prog["drgn_test_compound_page"]), 2 * self.prog["PAGE_SIZE"]
+        )
+
     @skip_unless_have_full_mm_support
     def test_decode_page_flags(self):
         with self._pages() as (map, _, pfns):

tests/linux_kernel/kmod/drgn_test.c

@@ -118,12 +118,16 @@ void *drgn_test_va;
 phys_addr_t drgn_test_pa;
 unsigned long drgn_test_pfn;
 struct page *drgn_test_page;
+struct page *drgn_test_compound_page;
 
 static int drgn_test_mm_init(void)
 {
 	drgn_test_page = alloc_page(GFP_KERNEL);
 	if (!drgn_test_page)
 		return -ENOMEM;
+	drgn_test_compound_page = alloc_pages(GFP_KERNEL | __GFP_COMP, 1);
+	if (!drgn_test_compound_page)
+		return -ENOMEM;
 	drgn_test_va = page_address(drgn_test_page);
 	drgn_test_pa = virt_to_phys(drgn_test_va);
 	drgn_test_pfn = PHYS_PFN(drgn_test_pa);
@@ -132,6 +136,8 @@ static int drgn_test_mm_init(void)
 
 static void drgn_test_mm_exit(void)
 {
+	if (drgn_test_compound_page)
+		__free_pages(drgn_test_page, 1);
 	if (drgn_test_page)
 		__free_pages(drgn_test_page, 0);
 }