# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later

"""
Memory Management
-----------------

The ``drgn.helpers.linux.mm`` module provides helpers for working with the
Linux memory management (MM) subsystem.

.. _virtual address translation failures:

Helpers that translate virtual addresses or read virtual memory may fail for
multiple reasons:

1. If the address is invalid.
2. If the address is swapped or paged out.
3. If the address is in `high memory
   <https://docs.kernel.org/mm/highmem.html>`_. High memory is only used for
   userspace memory by 32-bit systems with a lot of physical memory, and only
   if ``CONFIG_HIGHMEM`` is enabled.

   3a. If the page table is in high memory. This is only possible if
   ``CONFIG_HIGHPTE`` is enabled.
"""

import operator
import re
from typing import Callable, Iterator, List, Optional

from _drgn import (
    _linux_helper_direct_mapping_offset,
    _linux_helper_follow_phys,
    _linux_helper_read_vm,
)
from drgn import NULL, IntegerLike, Object, ObjectAbsentError, Program, cast
from drgn.helpers.common.format import decode_enum_type_flags
from drgn.helpers.common.prog import takes_program_or_default
from drgn.helpers.linux.list import list_for_each_entry
from drgn.helpers.linux.mapletree import mt_for_each, mtree_load
from drgn.helpers.linux.rbtree import rb_find

__all__ = (
    "PFN_PHYS",
    "PHYS_PFN",
    "PageCompound",
    "PageHead",
    "PageSlab",
    "PageTail",
    "access_process_vm",
    "access_remote_vm",
    "cmdline",
    "compound_head",
    "compound_nr",
    "compound_order",
    "decode_page_flags",
    "environ",
    "find_vmap_area",
    "follow_page",
    "follow_pfn",
    "follow_phys",
    "for_each_page",
    "for_each_vma",
    "for_each_vmap_area",
    "page_size",
    "page_to_pfn",
    "page_to_phys",
    "page_to_virt",
    "pfn_to_page",
    "pfn_to_virt",
    "phys_to_page",
    "phys_to_virt",
    "totalram_pages",
    "virt_to_page",
    "virt_to_pfn",
    "virt_to_phys",
    "vma_find",
    "vmalloc_to_page",
    "vmalloc_to_pfn",
    # Generated by scripts/generate_page_flag_getters.py.
    "PageActive",
    "PageChecked",
    "PageDirty",
    "PageDoubleMap",
    "PageError",
    "PageForeign",
    "PageHWPoison",
    "PageHasHWPoisoned",
    "PageIdle",
    "PageIsolated",
    "PageLRU",
    "PageLocked",
    "PageMappedToDisk",
    "PageMlocked",
    "PageOwnerPriv1",
    "PagePinned",
    "PagePrivate",
    "PagePrivate2",
    "PageReadahead",
    "PageReclaim",
    "PageReferenced",
    "PageReported",
    "PageReserved",
    "PageSavePinned",
    "PageSkipKASanPoison",
    "PageSlobFree",
    "PageSwapBacked",
    "PageUncached",
    "PageUnevictable",
    "PageUptodate",
    "PageVmemmapSelfHosted",
    "PageWaiters",
    "PageWorkingset",
    "PageWriteback",
    "PageXenRemapped",
    "PageYoung",
)


def PageActive(page: Object) -> bool:
    """
    Return whether the ``PG_active`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_active"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageChecked(page: Object) -> bool:
    """
    Return whether the ``PG_checked`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_checked"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageDirty(page: Object) -> bool:
    """
    Return whether the ``PG_dirty`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_dirty"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageDoubleMap(page: Object) -> bool:
    """
    Return whether the ``PG_double_map`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_double_map"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageError(page: Object) -> bool:
    """
    Return whether the ``PG_error`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_error"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageForeign(page: Object) -> bool:
    """
    Return whether the ``PG_foreign`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_foreign"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageHWPoison(page: Object) -> bool:
    """
    Return whether the ``PG_hwpoison`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_hwpoison"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageHasHWPoisoned(page: Object) -> bool:
    """
    Return whether the ``PG_has_hwpoisoned`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_has_hwpoisoned"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageIdle(page: Object) -> bool:
    """
    Return whether the ``PG_idle`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_idle"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageIsolated(page: Object) -> bool:
    """
    Return whether the ``PG_isolated`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_isolated"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageLRU(page: Object) -> bool:
    """
    Return whether the ``PG_lru`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_lru"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageLocked(page: Object) -> bool:
    """
    Return whether the ``PG_locked`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_locked"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageMappedToDisk(page: Object) -> bool:
    """
    Return whether the ``PG_mappedtodisk`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_mappedtodisk"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageMlocked(page: Object) -> bool:
    """
    Return whether the ``PG_mlocked`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_mlocked"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageOwnerPriv1(page: Object) -> bool:
    """
    Return whether the ``PG_owner_priv_1`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_owner_priv_1"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PagePinned(page: Object) -> bool:
    """
    Return whether the ``PG_pinned`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_pinned"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PagePrivate(page: Object) -> bool:
    """
    Return whether the ``PG_private`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_private"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PagePrivate2(page: Object) -> bool:
    """
    Return whether the ``PG_private_2`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_private_2"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageReadahead(page: Object) -> bool:
    """
    Return whether the ``PG_readahead`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_readahead"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageReclaim(page: Object) -> bool:
    """
    Return whether the ``PG_reclaim`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_reclaim"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageReferenced(page: Object) -> bool:
    """
    Return whether the ``PG_referenced`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_referenced"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageReported(page: Object) -> bool:
    """
    Return whether the ``PG_reported`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_reported"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageReserved(page: Object) -> bool:
    """
    Return whether the ``PG_reserved`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_reserved"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageSavePinned(page: Object) -> bool:
    """
    Return whether the ``PG_savepinned`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_savepinned"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageSkipKASanPoison(page: Object) -> bool:
    """
    Return whether the ``PG_skip_kasan_poison`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_skip_kasan_poison"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageSlobFree(page: Object) -> bool:
    """
    Return whether the ``PG_slob_free`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_slob_free"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageSwapBacked(page: Object) -> bool:
    """
    Return whether the ``PG_swapbacked`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_swapbacked"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageUncached(page: Object) -> bool:
    """
    Return whether the ``PG_uncached`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_uncached"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageUnevictable(page: Object) -> bool:
    """
    Return whether the ``PG_unevictable`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_unevictable"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageUptodate(page: Object) -> bool:
    """
    Return whether the ``PG_uptodate`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_uptodate"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageVmemmapSelfHosted(page: Object) -> bool:
    """
    Return whether the ``PG_vmemmap_self_hosted`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_vmemmap_self_hosted"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageWaiters(page: Object) -> bool:
    """
    Return whether the ``PG_waiters`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_waiters"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageWorkingset(page: Object) -> bool:
    """
    Return whether the ``PG_workingset`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_workingset"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageWriteback(page: Object) -> bool:
    """
    Return whether the ``PG_writeback`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_writeback"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageXenRemapped(page: Object) -> bool:
    """
    Return whether the ``PG_xen_remapped`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_xen_remapped"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


def PageYoung(page: Object) -> bool:
    """
    Return whether the ``PG_young`` flag is set on a page.

    :param page: ``struct page *``
    """
    try:
        flag = page.prog_["PG_young"]
    except KeyError:
        return False
    return bool(page.flags & (1 << flag))


# End generated by scripts/generate_page_flag_getters.py.


def _get_PageSlab_impl(prog: Program) -> Callable[[Object], bool]:
    # Since Linux kernel commit 46df8e73a4a3 ("mm: free up PG_slab") (in
    # v6.10), slab pages are identified by a page type, which is indicated by a
    # mapcount value matching a value in VMCOREINFO. Before that, they are
    # indicated by a page flag.
    try:
        return prog.cache["PageSlab"]
    except KeyError:
        pass
    vmcoreinfo = prog["VMCOREINFO"].string_()
    match = re.search(
        rb"^NUMBER\(PAGE_SLAB_MAPCOUNT_VALUE\)=(-?[0-9]+)$", vmcoreinfo, flags=re.M
    )
    if match:
        PAGE_SLAB_MAPCOUNT_VALUE = int(match.group(1))

        def PageSlab(page: Object) -> bool:
            return page._mapcount.counter.value_() == PAGE_SLAB_MAPCOUNT_VALUE

    else:
        mask = 1 << prog["PG_slab"]

        def PageSlab(page: Object) -> bool:
            return bool(page.flags & mask)

    prog.cache["PageSlab"] = PageSlab
    return PageSlab


def PageSlab(page: Object) -> bool:
    """
    Return whether a page belongs to the slab allocator.

    :param page: ``struct page *``
    """
    return _get_PageSlab_impl(page.prog_)(page)


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_()
    # Since Linux kernel commit 1d798ca3f164 ("mm: make compound_head()
    # robust") (in v4.4), PG_head is always defined, and a tail page has the
    # least significant bit of compound_head set. Before that, there is no
    # compound_head (and no fake head pages). Instead, if
    # CONFIG_PAGEFLAGS_EXTENDED=y, then PG_head and PG_tail are defined.
    # Otherwise, there is only PG_compound, and PG_reclaim is set for tail
    # pages and clear for head pages.
    try:
        PG_head = page.prog_["PG_head"]
    except KeyError:
        return bool(page.flags & (1 << page.prog_["PG_compound"]))
    else:
        flags = page.flags.read_()
        if flags & (1 << PG_head):
            return True
        try:
            return bool(page.compound_head.read_() & 1)
        except AttributeError:
            return bool(flags & (1 << page.prog_["PG_tail"]))


# 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_()
    # See PageCompound() re: Linux kernel commit 1d798ca3f164 ("mm: make
    # compound_head() robust") (in v4.4).
    try:
        PG_head = page.prog_["PG_head"]
    except KeyError:
        PG_compound = page.prog_["PG_compound"]
        PG_head_mask = 1 << PG_compound
        PG_head_tail_mask = PG_head_mask | (1 << page.prog_["PG_reclaim"])
        return (page.flags & PG_head_tail_mask) == PG_head_mask
    else:
        if not (page.flags & (1 << PG_head)):
            return False
        try:
            return not _page_is_fake_head(page)
        except AttributeError:
            return True


def PageTail(page: Object) -> bool:
    """
    Return whether a page is a tail page in a `compound page`_.

    :param page: ``struct page *``
    """
    page = page.read_()
    # See PageCompound() re: Linux kernel commit 1d798ca3f164 ("mm: make
    # compound_head() robust") (in v4.4).
    try:
        if page.compound_head.value_() & 1:
            return True
    except AttributeError:
        try:
            PG_tail = page.prog_["PG_tail"]
        except KeyError:
            PG_head_tail_mask = (1 << page.prog_["PG_compound"]) | (
                1 << page.prog_["PG_reclaim"]
            )
            return (page.flags & PG_head_tail_mask) == PG_head_tail_mask
        else:
            return bool(page.flags & (1 << PG_tail))
    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_()
    try:
        head = page.compound_head.read_()
    except AttributeError:
        # Before Linux kernel commit 1d798ca3f164 ("mm: make compound_head()
        # robust") (in v4.4), the head page is in page->first_page, and there
        # are no fake head pages.
        return page.first_page.read_() if PageTail(page) else page
    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``
    """
    prog = page.prog_

    if not PageHead(page):
        return Object(prog, "unsigned int", 0)

    # Since Linux kernel commit ebc1baf5c9b4 ("mm: free up a word in the first
    # tail page") (in v6.6), the compound order is in the low byte of struct
    # folio::_flags_1 (from_folio = 2). Between that and Linux kernel commit
    # Linux kernel commit 379708ffde1b ("mm: add the first tail page to struct
    # folio") (in v6.1), the compound order is in struct folio::_folio_order
    # (from_folio = 1). Before Linux kernel commit 1c5509be58f6 ("mm: remove
    # 'First tail page' members from struct page") (in v6.3), the compound
    # order is in struct page::compound_order of the first tail page
    # (from_folio = 0).
    try:
        from_folio = prog.cache["compound_order_from_folio"]
    except KeyError:
        from_folio = 0
        try:
            struct_folio = prog.type("struct folio")
        except LookupError:
            pass
        else:
            if struct_folio.has_member("_folio_order"):
                from_folio = 1
            elif struct_folio.has_member("_flags_1"):
                from_folio = 2
        prog.cache["compound_order_from_folio"] = from_folio
    if from_folio == 2:
        return cast("unsigned int", cast("struct folio *", page)._flags_1 & 0xFF)
    elif from_folio == 1:
        return cast("unsigned int", cast("struct folio *", page)._folio_order)
    else:
        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``
    """
    return Object(page.prog_, "unsigned long", 1) << compound_order(page)


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.

    >>> decode_page_flags(page)
    'PG_uptodate|PG_dirty|PG_lru|PG_reclaim|PG_swapbacked|PG_readahead|PG_savepinned|PG_isolated|PG_reported'

    :param page: ``struct page *``
    """
    NR_PAGEFLAGS = page.prog_["__NR_PAGEFLAGS"]
    PAGEFLAGS_MASK = (1 << NR_PAGEFLAGS.value_()) - 1
    return decode_enum_type_flags(
        page.flags.value_() & PAGEFLAGS_MASK, NR_PAGEFLAGS.type_
    )


# Get the struct page * for PFN 0.
def _page0(prog: Program) -> Object:
    try:
        return prog.cache["page0"]
    except KeyError:
        pass
    try:
        # With CONFIG_SPARSEMEM_VMEMMAP=y, page 0 is vmemmap.
        page0 = prog["vmemmap"]
    except KeyError:
        contig_page_data = prog["contig_page_data"]
        # With CONFIG_FLATMEM=y, page 0 is mem_map - ARCH_PFN_OFFSET, but we
        # can't determine ARCH_PFN_OFFSET easily. Alternatively,
        # contig_page_data.node_mem_map is the struct page * for
        # contig_page_data.node_start_pfn, therefore page 0 is:
        page0 = contig_page_data.node_mem_map - contig_page_data.node_start_pfn
    # The struct page array is not contiguous for CONFIG_SPARSEMEM=y with
    # CONFIG_SPARSEMEM_VMEMMAP=n or CONFIG_DISCONTIGMEM=y, so those are not
    # supported yet.
    prog.cache["page0"] = page0
    return page0


@takes_program_or_default
def for_each_page(prog: Program) -> Iterator[Object]:
    """
    Iterate over every ``struct page *`` from the minimum to the maximum page.

    .. note::

        This may include offline pages which don't have a valid ``struct
        page``. Wrap accesses in a ``try`` ... ``except``
        :class:`drgn.FaultError`:

        >>> for page in for_each_page():
        ...     try:
        ...         if PageLRU(page):
        ...             print(hex(page))
        ...     except drgn.FaultError:
        ...         continue
        0xfffffb4a000c0000
        0xfffffb4a000c0040
        ...

        This may be fixed in the future.

    :return: Iterator of ``struct page *`` objects.
    """
    page0 = _page0(prog)
    for i in range(prog["min_low_pfn"], prog["max_pfn"]):
        yield page0 + i


@takes_program_or_default
def PFN_PHYS(prog: Program, pfn: IntegerLike) -> Object:
    """
    Get the physical address of a page frame number (PFN).

    :param pfn: ``unsigned long``
    :return: ``phys_addr_t``
    """
    return Object(prog, "phys_addr_t", pfn) << prog["PAGE_SHIFT"]


@takes_program_or_default
def PHYS_PFN(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the page frame number (PFN) of a physical address.

    :param addr: ``phys_addr_t``
    :return: ``unsigned long``
    """
    return Object(prog, "unsigned long", addr) >> prog["PAGE_SHIFT"]


def page_to_pfn(page: Object) -> Object:
    """
    Get the page frame number (PFN) of a page.

    :param page: ``struct page *``
    :return: ``unsigned long``
    """
    return cast("unsigned long", page - _page0(page.prog_))


def page_to_phys(page: Object) -> Object:
    """
    Get the physical address of a page.

    :param page: ``struct page *``
    :return: ``phys_addr_t``
    """
    return PFN_PHYS(page_to_pfn(page))


def page_to_virt(page: Object) -> Object:
    """
    Get the directly mapped virtual address of a page.

    :param page: ``struct page *``
    :return: ``void *``
    """
    return pfn_to_virt(page_to_pfn(page))


@takes_program_or_default
def pfn_to_page(prog: Program, pfn: IntegerLike) -> Object:
    """
    Get the page with a page frame number (PFN).

    :param pfn: ``unsigned long``
    :return: ``struct page *``
    """
    return _page0(prog) + pfn


@takes_program_or_default
def pfn_to_virt(prog: Program, pfn: IntegerLike) -> Object:
    """
    Get the directly mapped virtual address of a page frame number (PFN).

    :param pfn: ``unsigned long``
    :return: ``void *``
    """
    return phys_to_virt(PFN_PHYS(prog, pfn))


@takes_program_or_default
def phys_to_page(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the page containing a physical address.

    :param addr: ``phys_addr_t``
    :return: ``struct page *``
    """
    return pfn_to_page(PHYS_PFN(prog, addr))


@takes_program_or_default
def phys_to_virt(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the directly mapped virtual address of a physical address.

    :param addr: ``phys_addr_t``
    :return: ``void *``
    """
    return Object(
        prog, "void *", operator.index(addr) + _linux_helper_direct_mapping_offset(prog)
    )


@takes_program_or_default
def virt_to_page(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the page containing a directly mapped virtual address.

    .. _mm-helpers-direct-map:

    .. note::

        This only works for virtual addresses from the "direct map". This
        includes address from:

        * kmalloc
        * Slab allocator
        * Page allocator

        But not:

        * vmalloc
        * vmap
        * ioremap
        * Symbols (function pointers, global variables)

        For vmalloc or vmap addresses, use :func:`vmalloc_to_page(addr)
        <vmalloc_to_page>`. For arbitrary kernel addresses, use
        :func:`follow_page(prog["init_mm"].address_of_(), addr) <follow_page>`.

    :param addr: ``void *``
    :return: ``struct page *``
    """
    return pfn_to_page(virt_to_pfn(prog, addr))


@takes_program_or_default
def virt_to_pfn(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the page frame number (PFN) of a directly mapped virtual address.

    .. note::

        This only works for virtual addresses from the :ref:`"direct map"
        <mm-helpers-direct-map>`. For vmalloc or vmap addresses, use
        :func:`vmalloc_to_pfn(addr) <vmalloc_to_pfn>`. For arbitrary kernel
        addresses, use :func:`follow_pfn(prog["init_mm"].address_of_(), addr)
        <follow_pfn>`.

    :param addr: ``void *``
    :return: ``unsigned long``
    """
    return PHYS_PFN(virt_to_phys(prog, addr))


@takes_program_or_default
def virt_to_phys(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the physical address of a directly mapped virtual address.

    .. note::

        This only works for virtual addresses from the :ref:`"direct map"
        <mm-helpers-direct-map>`. For arbitrary kernel addresses, use
        :func:`follow_phys(prog["init_mm"].address_of_(), addr) <follow_phys>`.

    :param addr: ``void *``
    :return: ``phys_addr_t``
    """
    return Object(
        prog,
        "unsigned long",
        operator.index(addr) - _linux_helper_direct_mapping_offset(prog),
    )


def follow_page(mm: Object, addr: IntegerLike) -> Object:
    """
    Get the page that a virtual address maps to in a virtual address space.

    >>> task = find_task(113)
    >>> follow_page(task.mm, 0x7fffbbb6d4d0)
    *(struct page *)0xffffbe4bc0337b80 = {
        ...
    }

    :param mm: ``struct mm_struct *``
    :param addr: ``void *``
    :return: ``struct page *``
    :raises FaultError: if the virtual address
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    return phys_to_page(follow_phys(mm, addr))


def follow_pfn(mm: Object, addr: IntegerLike) -> Object:
    """
    Get the page frame number (PFN) that a virtual address maps to in a virtual
    address space.

    >>> task = find_task(113)
    >>> follow_pfn(task.mm, 0x7fffbbb6d4d0)
    (unsigned long)52718

    :param mm: ``struct mm_struct *``
    :param addr: ``void *``
    :return: ``unsigned long``
    :raises FaultError: if the virtual address
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    return PHYS_PFN(follow_phys(mm, addr))


def follow_phys(mm: Object, addr: IntegerLike) -> Object:
    """
    Get the physical address that a virtual address maps to in a virtual
    address space.

    >>> task = find_task(113)
    >>> follow_phys(task.mm, 0x7fffbbb6d4d0)
    (phys_addr_t)215934160

    :param mm: ``struct mm_struct *``
    :param addr: ``void *``
    :return: ``phys_addr_t``
    :raises FaultError: if the virtual address
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    prog = mm.prog_
    return Object(prog, "phys_addr_t", _linux_helper_follow_phys(prog, mm.pgd, addr))


@takes_program_or_default
def vmalloc_to_page(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the page containing a vmalloc or vmap address.

    >>> task = find_task(113)
    >>> vmalloc_to_page(task.stack)
    *(struct page *)0xffffbe4bc00a2200 = {
        ...
    }

    :param addr: ``void *``
    :return: ``struct page *``
    """
    return follow_page(prog["init_mm"].address_of_(), addr)


@takes_program_or_default
def vmalloc_to_pfn(prog: Program, addr: IntegerLike) -> Object:
    """
    Get the page frame number (PFN) containing a vmalloc or vmap address.

    >>> task = find_task(113)
    >>> vmalloc_to_pfn(task.stack)
    (unsigned long)10376

    :param addr: ``void *``
    :return: ``unsigned long``
    """
    return page_to_pfn(vmalloc_to_page(prog, addr))


def _vmap_area_rb_cmp(addr: int, va: Object) -> int:
    if addr < va.va_start.value_():
        return -1
    elif addr >= va.va_end.value_():
        return 1
    else:
        return 0


def _vmap_nodes(prog: Program) -> Object:
    vmap_nodes = prog["vmap_nodes"]
    try:
        return vmap_nodes.read_()
    except ObjectAbsentError:
        # On !SMP and 32-bit kernels, vmap_nodes is initialized to &single and
        # never reassigned. GCC as of version 12.2 doesn't generate a location
        # for vmap_nodes description in that case.
        return prog.variable("single", "mm/vmalloc.c").address_of_()


@takes_program_or_default
def find_vmap_area(prog: Program, addr: IntegerLike) -> Object:
    """
    Return the ``struct vmap_area *`` containing an address.

    >>> find_vmap_area(0xffffa2b680081000)
    *(struct vmap_area *)0xffffa16541046b40 = {
            ...
    }

    :param addr: Address to look up.
    :return: ``struct vmap_area *`` (``NULL`` if not found)
    """
    addr = operator.index(addr)
    # Since Linux kernel commit d093602919ad ("mm: vmalloc: remove global
    # vmap_area_root rb-tree") (in v6.9), vmap areas are split up in multiple
    # red-black trees in separate "nodes". Before that, they're in a single
    # red-black tree.
    try:
        vmap_nodes = _vmap_nodes(prog)
    except KeyError:
        return rb_find(
            "struct vmap_area",
            prog["vmap_area_root"].address_of_(),
            "rb_node",
            addr,
            _vmap_area_rb_cmp,
        )
    else:
        nr_vmap_nodes = prog["nr_vmap_nodes"].value_()
        i = j = (addr // prog["vmap_zone_size"].value_()) % nr_vmap_nodes
        while True:
            vn = vmap_nodes[i]
            va = rb_find(
                "struct vmap_area",
                vn.busy.root.address_of_(),
                "rb_node",
                addr,
                _vmap_area_rb_cmp,
            )
            if va:
                return va
            # As noted in the kernel implementation, the given address may be
            # in a different node than the start address, so we have to loop.
            i = (i + 1) % nr_vmap_nodes
            if i == j:
                return NULL(prog, "struct vmap_area *")


@takes_program_or_default
def for_each_vmap_area(prog: Program) -> Iterator[Object]:
    """
    Iterate over every ``struct vmap_area *`` on the system.

    >>> for va in for_each_vmap_area():
    ...     caller = ""
    ...     if va.vm:
    ...         try:
    ...             sym = prog.symbol(va.vm.caller)
    ...         except LookupError:
    ...             pass
    ...         else:
    ...             caller = f" {sym.name}"
    ...     print(f"{hex(va.va_start)}-{hex(va.va_end)}{caller}")
    ...
    0xffffa2b680000000-0xffffa2b680005000 irq_init_percpu_irqstack
    0xffffa2b680005000-0xffffa2b680007000 acpi_os_map_iomem
    0xffffa2b68000b000-0xffffa2b68000d000 hpet_enable
    0xffffa2b680080000-0xffffa2b680085000 kernel_clone
    ...

    :return: Iterator of ``struct vmap_area *`` objects.
    """
    # Since Linux kernel commit d093602919ad ("mm: vmalloc: remove global
    # vmap_area_root rb-tree") (in v6.9), vmap areas are split up in multiple
    # lists in separate "nodes". Before that, they're in a single list.
    try:
        vmap_nodes = _vmap_nodes(prog)
    except KeyError:
        yield from list_for_each_entry(
            "struct vmap_area", prog["vmap_area_list"].address_of_(), "list"
        )
    else:
        for i in range(prog["nr_vmap_nodes"]):
            yield from list_for_each_entry(
                "struct vmap_area", vmap_nodes[i].busy.head.address_of_(), "list"
            )


def access_process_vm(task: Object, address: IntegerLike, size: IntegerLike) -> bytes:
    """
    Read memory from a task's virtual address space.

    >>> task = find_task(1490152)
    >>> access_process_vm(task, 0x7f8a62b56da0, 12)
    b'hello, world'

    :param task: ``struct task_struct *``
    :param address: Starting address.
    :param size: Number of bytes to read.
    :raises FaultError: if the virtual address
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    return _linux_helper_read_vm(task.prog_, task.mm.pgd, address, size)


def access_remote_vm(mm: Object, address: IntegerLike, size: IntegerLike) -> bytes:
    """
    Read memory from a virtual address space. This is similar to
    :func:`access_process_vm()`, but it takes a ``struct mm_struct *`` instead
    of a ``struct task_struct *``.

    >>> task = find_task(1490152)
    >>> access_remote_vm(task.mm, 0x7f8a62b56da0, 12)
    b'hello, world'

    :param mm: ``struct mm_struct *``
    :param address: Starting address.
    :param size: Number of bytes to read.
    :raises FaultError: if the virtual address
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    return _linux_helper_read_vm(mm.prog_, mm.pgd, address, size)


def cmdline(task: Object) -> Optional[List[bytes]]:
    """
    Get the list of command line arguments of a task, or ``None`` for kernel tasks.

    >>> cmdline(find_task(1495216))
    [b'vim', b'drgn/helpers/linux/mm.py']

    .. code-block:: console

        $ tr '\\0' ' ' < /proc/1495216/cmdline
        vim drgn/helpers/linux/mm.py

    :param task: ``struct task_struct *``
    :raises FaultError: if the virtual address containing the command line
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    mm = task.mm.read_()
    if not mm:
        return None
    arg_start = mm.arg_start.value_()
    arg_end = mm.arg_end.value_()
    return access_remote_vm(mm, arg_start, arg_end - arg_start).split(b"\0")[:-1]


def environ(task: Object) -> Optional[List[bytes]]:
    """
    Get the list of environment variables of a task, or ``None`` for kernel tasks.

    >>> environ(find_task(1497797))
    [b'HOME=/root', b'PATH=/usr/local/sbin:/usr/local/bin:/usr/bin', b'LOGNAME=root']

    .. code-block:: console

        $ tr '\\0' '\\n' < /proc/1497797/environ
        HOME=/root
        PATH=/usr/local/sbin:/usr/local/bin:/usr/bin
        LOGNAME=root

    :param task: ``struct task_struct *``
    :raises FaultError: if the virtual address containing the environment
        :ref:`cannot be translated <virtual address translation failures>`
    :raises NotImplementedError: if virtual address translation is :ref:`not
        supported <architecture support matrix>` for this architecture yet
    """
    mm = task.mm.read_()
    if not mm:
        return None
    env_start = mm.env_start.value_()
    env_end = mm.env_end.value_()
    return access_remote_vm(mm, env_start, env_end - env_start).split(b"\0")[:-1]


def _vma_rb_cmp(addr: int, vma: Object) -> int:
    if addr < vma.vm_start.value_():
        return -1
    elif addr >= vma.vm_end.value_():
        return 1
    else:
        return 0


def vma_find(mm: Object, addr: IntegerLike) -> Object:
    """
    Return the virtual memory area (VMA) containing an address.

    :param mm: ``struct mm_struct *``
    :param addr: Address to look up.
    :return: ``struct vm_area_struct *`` (``NULL`` if not found)
    """
    try:
        # Since Linux kernel commit 524e00b36e8c ("mm: remove rb tree.") (in
        # v6.1), VMAs are stored in a maple tree.
        mt = mm.mm_mt.address_of_()
    except AttributeError:
        # Before that, they are in a red-black tree.
        return rb_find(
            "struct vm_area_struct",
            mm.mm_rb.address_of_(),
            "vm_rb",
            operator.index(addr),
            _vma_rb_cmp,
        )
    else:
        return cast("struct vm_area_struct *", mtree_load(mt, addr))


def for_each_vma(mm: Object) -> Iterator[Object]:
    """
    Iterate over every virtual memory area (VMA) in a virtual address space.

    >>> for vma in for_each_vma(task.mm):
    ...     print(vma)
    ...
    *(struct vm_area_struct *)0xffff97ad82bfc930 = {
        ...
    }
    *(struct vm_area_struct *)0xffff97ad82bfc0a8 = {
        ...
    }
    ...

    :param mm: ``struct mm_struct *``
    :return: Iterator of ``struct vm_area_struct *`` objects.
    """
    try:
        # Since Linux kernel commit 763ecb035029 ("mm: remove the vma linked
        # list") (in v6.1), VMAs are stored in a maple tree.
        mt = mm.mm_mt.address_of_()
    except AttributeError:
        # Before that, they are in a linked list.
        vma = mm.mmap
        while vma:
            yield vma
            vma = vma.vm_next
    else:
        type = mm.prog_.type("struct vm_area_struct *")
        for _, _, entry in mt_for_each(mt):
            yield cast(type, entry)


@takes_program_or_default
def totalram_pages(prog: Program) -> int:
    """Return the total number of RAM pages."""
    try:
        # The variable is present since Linux kernel commit ca79b0c211af63fa32
        # ("mm: convert totalram_pages and totalhigh_pages variables
        # to atomic") (in v5.0).
        return prog["_totalram_pages"].counter.value_()
    except KeyError:
        return prog["totalram_pages"].value_()


@takes_program_or_default
def in_direct_map(prog: Program, addr: IntegerLike) -> bool:
    """
    Return whether an address is within the kernel's direct memory mapping.

    :param addr: address to check
    """
    addr = operator.index(addr)
    start_addr = pfn_to_virt(prog["min_low_pfn"]).value_()
    end_addr = (pfn_to_virt(prog["max_low_pfn"]) + prog["PAGE_SIZE"]).value_()
    return start_addr <= addr < end_addr
