/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES
 *
 * Test the format API directly.
 *
 */
#include "kunit_iommu.h"
#include "pt_iter.h"

static void do_map(struct kunit *test, pt_vaddr_t va, pt_oaddr_t pa,
                   pt_vaddr_t len)
{
        struct kunit_iommu_priv *priv = test->priv;
        int ret;

        KUNIT_ASSERT_EQ(test, len, (size_t)len);

        ret = iommu_map(&priv->domain, va, pa, len, IOMMU_READ | IOMMU_WRITE,
                        GFP_KERNEL);
        KUNIT_ASSERT_NO_ERRNO_FN(test, "map_pages", ret);
}

#define KUNIT_ASSERT_PT_LOAD(test, pts, entry)             \
        ({                                                 \
                pt_load_entry(pts);                        \
                KUNIT_ASSERT_EQ(test, (pts)->type, entry); \
        })

struct check_levels_arg {
        struct kunit *test;
        void *fn_arg;
        void (*fn)(struct kunit *test, struct pt_state *pts, void *arg);
};

static int __check_all_levels(struct pt_range *range, void *arg,
                              unsigned int level, struct pt_table_p *table)
{
        struct pt_state pts = pt_init(range, level, table);
        struct check_levels_arg *chk = arg;
        struct kunit *test = chk->test;
        int ret;

        _pt_iter_first(&pts);


        /*
         * If we were able to use the full VA space this should always be the
         * last index in each table.
         */
        if (!(IS_32BIT && range->max_vasz_lg2 > 32)) {
                if (pt_feature(range->common, PT_FEAT_SIGN_EXTEND) &&
                    pts.level == pts.range->top_level)
                        KUNIT_ASSERT_EQ(test, pts.index,
                                        log2_to_int(range->max_vasz_lg2 - 1 -
                                                    pt_table_item_lg2sz(&pts)) -
                                                1);
                else
                        KUNIT_ASSERT_EQ(test, pts.index,
                                        log2_to_int(pt_table_oa_lg2sz(&pts) -
                                                    pt_table_item_lg2sz(&pts)) -
                                                1);
        }

        if (pt_can_have_table(&pts)) {
                pt_load_single_entry(&pts);
                KUNIT_ASSERT_EQ(test, pts.type, PT_ENTRY_TABLE);
                ret = pt_descend(&pts, arg, __check_all_levels);
                KUNIT_ASSERT_EQ(test, ret, 0);

                /* Index 0 is used by the test */
                if (IS_32BIT && !pts.index)
                        return 0;
                KUNIT_ASSERT_NE(chk->test, pts.index, 0);
        }

        /*
         * A format should not create a table with only one entry, at least this
         * test approach won't work.
         */
        KUNIT_ASSERT_GT(chk->test, pts.end_index, 1);

        /*
         * For increase top we end up using index 0 for the original top's tree,
         * so use index 1 for testing instead.
         */
        pts.index = 0;
        pt_index_to_va(&pts);
        pt_load_single_entry(&pts);
        if (pts.type == PT_ENTRY_TABLE && pts.end_index > 2) {
                pts.index = 1;
                pt_index_to_va(&pts);
        }
        (*chk->fn)(chk->test, &pts, chk->fn_arg);
        return 0;
}

/*
 * Call fn for each level in the table with a pts setup to index 0 in a table
 * for that level. This allows writing tests that run on every level.
 * The test can use every index in the table except the last one.
 */
static void check_all_levels(struct kunit *test,
                             void (*fn)(struct kunit *test,
                                        struct pt_state *pts, void *arg),
                             void *fn_arg)
{
        struct kunit_iommu_priv *priv = test->priv;
        struct pt_range range = pt_top_range(priv->common);
        struct check_levels_arg chk = {
                .test = test,
                .fn = fn,
                .fn_arg = fn_arg,
        };
        int ret;

        if (pt_feature(priv->common, PT_FEAT_DYNAMIC_TOP) &&
            priv->common->max_vasz_lg2 > range.max_vasz_lg2)
                range.last_va = fvalog2_set_mod_max(range.va,
                                                    priv->common->max_vasz_lg2);

        /*
         * Map a page at the highest VA, this will populate all the levels so we
         * can then iterate over them. Index 0 will be used for testing.
         */
        if (IS_32BIT && range.max_vasz_lg2 > 32)
                range.last_va = (u32)range.last_va;
        range.va = range.last_va - (priv->smallest_pgsz - 1);
        do_map(test, range.va, 0, priv->smallest_pgsz);

        range = pt_make_range(priv->common, range.va, range.last_va);
        ret = pt_walk_range(&range, __check_all_levels, &chk);
        KUNIT_ASSERT_EQ(test, ret, 0);
}

static void test_init(struct kunit *test)
{
        struct kunit_iommu_priv *priv = test->priv;

        /* Fixture does the setup */
        KUNIT_ASSERT_NE(test, priv->info.pgsize_bitmap, 0);
}

/*
 * Basic check that the log2_* functions are working, especially at the integer
 * limits.
 */
static void test_bitops(struct kunit *test)
{
        int i;

        KUNIT_ASSERT_EQ(test, fls_t(u32, 0), 0);
        KUNIT_ASSERT_EQ(test, fls_t(u32, 1), 1);
        KUNIT_ASSERT_EQ(test, fls_t(u32, BIT(2)), 3);
        KUNIT_ASSERT_EQ(test, fls_t(u32, U32_MAX), 32);

        KUNIT_ASSERT_EQ(test, fls_t(u64, 0), 0);
        KUNIT_ASSERT_EQ(test, fls_t(u64, 1), 1);
        KUNIT_ASSERT_EQ(test, fls_t(u64, BIT(2)), 3);
        KUNIT_ASSERT_EQ(test, fls_t(u64, U64_MAX), 64);

        KUNIT_ASSERT_EQ(test, ffs_t(u32, 1), 0);
        KUNIT_ASSERT_EQ(test, ffs_t(u32, BIT(2)), 2);
        KUNIT_ASSERT_EQ(test, ffs_t(u32, BIT(31)), 31);

        KUNIT_ASSERT_EQ(test, ffs_t(u64, 1), 0);
        KUNIT_ASSERT_EQ(test, ffs_t(u64, BIT(2)), 2);
        KUNIT_ASSERT_EQ(test, ffs_t(u64, BIT_ULL(63)), 63);

        for (i = 0; i != 31; i++)
                KUNIT_ASSERT_EQ(test, ffz_t(u64, BIT_ULL(i) - 1), i);

        for (i = 0; i != 63; i++)
                KUNIT_ASSERT_EQ(test, ffz_t(u64, BIT_ULL(i) - 1), i);

        for (i = 0; i != 32; i++) {
                u64 val = get_random_u64();

                KUNIT_ASSERT_EQ(test, log2_mod_t(u32, val, ffs_t(u32, val)), 0);
                KUNIT_ASSERT_EQ(test, log2_mod_t(u64, val, ffs_t(u64, val)), 0);

                KUNIT_ASSERT_EQ(test, log2_mod_t(u32, val, ffz_t(u32, val)),
                                log2_to_max_int_t(u32, ffz_t(u32, val)));
                KUNIT_ASSERT_EQ(test, log2_mod_t(u64, val, ffz_t(u64, val)),
                                log2_to_max_int_t(u64, ffz_t(u64, val)));
        }
}

static unsigned int ref_best_pgsize(pt_vaddr_t pgsz_bitmap, pt_vaddr_t va,
                                    pt_vaddr_t last_va, pt_oaddr_t oa)
{
        pt_vaddr_t pgsz_lg2;

        /* Brute force the constraints described in pt_compute_best_pgsize() */
        for (pgsz_lg2 = PT_VADDR_MAX_LG2 - 1; pgsz_lg2 != 0; pgsz_lg2--) {
                if ((pgsz_bitmap & log2_to_int(pgsz_lg2)) &&
                    log2_mod(va, pgsz_lg2) == 0 &&
                    oalog2_mod(oa, pgsz_lg2) == 0 &&
                    va + log2_to_int(pgsz_lg2) - 1 <= last_va &&
                    log2_div_eq(va, va + log2_to_int(pgsz_lg2) - 1, pgsz_lg2) &&
                    oalog2_div_eq(oa, oa + log2_to_int(pgsz_lg2) - 1, pgsz_lg2))
                        return pgsz_lg2;
        }
        return 0;
}

/* Check that the bit logic in pt_compute_best_pgsize() works. */
static void test_best_pgsize(struct kunit *test)
{
        unsigned int a_lg2;
        unsigned int b_lg2;
        unsigned int c_lg2;

        /* Try random prefixes with every suffix combination */
        for (a_lg2 = 1; a_lg2 != 10; a_lg2++) {
                for (b_lg2 = 1; b_lg2 != 10; b_lg2++) {
                        for (c_lg2 = 1; c_lg2 != 10; c_lg2++) {
                                pt_vaddr_t pgsz_bitmap = get_random_u64();
                                pt_vaddr_t va = get_random_u64() << a_lg2;
                                pt_oaddr_t oa = get_random_u64() << b_lg2;
                                pt_vaddr_t last_va = log2_set_mod_max(
                                        get_random_u64(), c_lg2);

                                if (va > last_va)
                                        swap(va, last_va);
                                KUNIT_ASSERT_EQ(
                                        test,
                                        pt_compute_best_pgsize(pgsz_bitmap, va,
                                                               last_va, oa),
                                        ref_best_pgsize(pgsz_bitmap, va,
                                                        last_va, oa));
                        }
                }
        }

        /* 0 prefix, every suffix */
        for (c_lg2 = 1; c_lg2 != PT_VADDR_MAX_LG2 - 1; c_lg2++) {
                pt_vaddr_t pgsz_bitmap = get_random_u64();
                pt_vaddr_t va = 0;
                pt_oaddr_t oa = 0;
                pt_vaddr_t last_va = log2_set_mod_max(0, c_lg2);

                KUNIT_ASSERT_EQ(test,
                                pt_compute_best_pgsize(pgsz_bitmap, va, last_va,
                                                       oa),
                                ref_best_pgsize(pgsz_bitmap, va, last_va, oa));
        }

        /* 1's prefix, every suffix */
        for (a_lg2 = 1; a_lg2 != 10; a_lg2++) {
                for (b_lg2 = 1; b_lg2 != 10; b_lg2++) {
                        for (c_lg2 = 1; c_lg2 != 10; c_lg2++) {
                                pt_vaddr_t pgsz_bitmap = get_random_u64();
                                pt_vaddr_t va = PT_VADDR_MAX << a_lg2;
                                pt_oaddr_t oa = PT_VADDR_MAX << b_lg2;
                                pt_vaddr_t last_va = PT_VADDR_MAX;

                                KUNIT_ASSERT_EQ(
                                        test,
                                        pt_compute_best_pgsize(pgsz_bitmap, va,
                                                               last_va, oa),
                                        ref_best_pgsize(pgsz_bitmap, va,
                                                        last_va, oa));
                        }
                }
        }

        /* pgsize_bitmap is always 0 */
        for (a_lg2 = 1; a_lg2 != 10; a_lg2++) {
                for (b_lg2 = 1; b_lg2 != 10; b_lg2++) {
                        for (c_lg2 = 1; c_lg2 != 10; c_lg2++) {
                                pt_vaddr_t pgsz_bitmap = 0;
                                pt_vaddr_t va = get_random_u64() << a_lg2;
                                pt_oaddr_t oa = get_random_u64() << b_lg2;
                                pt_vaddr_t last_va = log2_set_mod_max(
                                        get_random_u64(), c_lg2);

                                if (va > last_va)
                                        swap(va, last_va);
                                KUNIT_ASSERT_EQ(
                                        test,
                                        pt_compute_best_pgsize(pgsz_bitmap, va,
                                                               last_va, oa),
                                        0);
                        }
                }
        }

        if (sizeof(pt_vaddr_t) <= 4)
                return;

        /* over 32 bit page sizes */
        for (a_lg2 = 32; a_lg2 != 42; a_lg2++) {
                for (b_lg2 = 32; b_lg2 != 42; b_lg2++) {
                        for (c_lg2 = 32; c_lg2 != 42; c_lg2++) {
                                pt_vaddr_t pgsz_bitmap = get_random_u64();
                                pt_vaddr_t va = get_random_u64() << a_lg2;
                                pt_oaddr_t oa = get_random_u64() << b_lg2;
                                pt_vaddr_t last_va = log2_set_mod_max(
                                        get_random_u64(), c_lg2);

                                if (va > last_va)
                                        swap(va, last_va);
                                KUNIT_ASSERT_EQ(
                                        test,
                                        pt_compute_best_pgsize(pgsz_bitmap, va,
                                                               last_va, oa),
                                        ref_best_pgsize(pgsz_bitmap, va,
                                                        last_va, oa));
                        }
                }
        }
}

/*
 * Check that pt_install_table() and pt_table_pa() match
 */
static void test_lvl_table_ptr(struct kunit *test, struct pt_state *pts,
                               void *arg)
{
        struct kunit_iommu_priv *priv = test->priv;
        pt_oaddr_t paddr =
                log2_set_mod(priv->test_oa, 0, priv->smallest_pgsz_lg2);
        struct pt_write_attrs attrs = {};

        if (!pt_can_have_table(pts))
                return;

        KUNIT_ASSERT_NO_ERRNO_FN(test, "pt_iommu_set_prot",
                                 pt_iommu_set_prot(pts->range->common, &attrs,
                                                   IOMMU_READ));

        pt_load_single_entry(pts);
        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_EMPTY);

        KUNIT_ASSERT_TRUE(test, pt_install_table(pts, paddr, &attrs));

        /* A second install should pass because install updates pts->entry. */
        KUNIT_ASSERT_EQ(test, pt_install_table(pts, paddr, &attrs), true);

        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_TABLE);
        KUNIT_ASSERT_EQ(test, pt_table_pa(pts), paddr);

        pt_clear_entries(pts, ilog2(1));
        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_EMPTY);
}

static void test_table_ptr(struct kunit *test)
{
        check_all_levels(test, test_lvl_table_ptr, NULL);
}

struct lvl_radix_arg {
        pt_vaddr_t vbits;
};

/*
 * Check pt_table_oa_lg2sz() and pt_table_item_lg2sz() they need to decode a
 * continuous list of VA across all the levels that covers the entire advertised
 * VA space.
 */
static void test_lvl_radix(struct kunit *test, struct pt_state *pts, void *arg)
{
        unsigned int table_lg2sz = pt_table_oa_lg2sz(pts);
        unsigned int isz_lg2 = pt_table_item_lg2sz(pts);
        struct lvl_radix_arg *radix = arg;

        /* Every bit below us is decoded */
        KUNIT_ASSERT_EQ(test, log2_set_mod_max(0, isz_lg2), radix->vbits);

        /* We are not decoding bits someone else is */
        KUNIT_ASSERT_EQ(test, log2_div(radix->vbits, isz_lg2), 0);

        /* Can't decode past the pt_vaddr_t size */
        KUNIT_ASSERT_LE(test, table_lg2sz, PT_VADDR_MAX_LG2);
        KUNIT_ASSERT_EQ(test, fvalog2_div(table_lg2sz, PT_MAX_VA_ADDRESS_LG2),
                        0);

        radix->vbits = fvalog2_set_mod_max(0, table_lg2sz);
}

static void test_max_va(struct kunit *test)
{
        struct kunit_iommu_priv *priv = test->priv;
        struct pt_range range = pt_top_range(priv->common);

        KUNIT_ASSERT_GE(test, priv->common->max_vasz_lg2, range.max_vasz_lg2);
}

static void test_table_radix(struct kunit *test)
{
        struct kunit_iommu_priv *priv = test->priv;
        struct lvl_radix_arg radix = { .vbits = priv->smallest_pgsz - 1 };
        struct pt_range range;

        check_all_levels(test, test_lvl_radix, &radix);

        range = pt_top_range(priv->common);
        if (range.max_vasz_lg2 == PT_VADDR_MAX_LG2) {
                KUNIT_ASSERT_EQ(test, radix.vbits, PT_VADDR_MAX);
        } else {
                if (!IS_32BIT)
                        KUNIT_ASSERT_EQ(test,
                                        log2_set_mod_max(0, range.max_vasz_lg2),
                                        radix.vbits);
                KUNIT_ASSERT_EQ(test, log2_div(radix.vbits, range.max_vasz_lg2),
                                0);
        }
}

static unsigned int safe_pt_num_items_lg2(const struct pt_state *pts)
{
        struct pt_range top_range = pt_top_range(pts->range->common);
        struct pt_state top_pts = pt_init_top(&top_range);

        /*
         * Avoid calling pt_num_items_lg2() on the top, instead we can derive
         * the size of the top table from the top range.
         */
        if (pts->level == top_range.top_level)
                return ilog2(pt_range_to_end_index(&top_pts));
        return pt_num_items_lg2(pts);
}

static void test_lvl_possible_sizes(struct kunit *test, struct pt_state *pts,
                                    void *arg)
{
        unsigned int num_items_lg2 = safe_pt_num_items_lg2(pts);
        pt_vaddr_t pgsize_bitmap = pt_possible_sizes(pts);
        unsigned int isz_lg2 = pt_table_item_lg2sz(pts);

        if (!pt_can_have_leaf(pts)) {
                KUNIT_ASSERT_EQ(test, pgsize_bitmap, 0);
                return;
        }

        /* No bits for sizes that would be outside this table */
        KUNIT_ASSERT_EQ(test, log2_mod(pgsize_bitmap, isz_lg2), 0);
        KUNIT_ASSERT_EQ(
                test, fvalog2_div(pgsize_bitmap, num_items_lg2 + isz_lg2), 0);

        /*
         * Non contiguous must be supported. AMDv1 has a HW bug where it does
         * not support it on one of the levels.
         */
        if ((u64)pgsize_bitmap != 0xff0000000000ULL ||
            strcmp(__stringify(PTPFX_RAW), "amdv1") != 0)
                KUNIT_ASSERT_TRUE(test, pgsize_bitmap & log2_to_int(isz_lg2));
        else
                KUNIT_ASSERT_NE(test, pgsize_bitmap, 0);

        /* A contiguous entry should not span the whole table */
        if (num_items_lg2 + isz_lg2 != PT_VADDR_MAX_LG2)
                KUNIT_ASSERT_FALSE(
                        test,
                        pgsize_bitmap & log2_to_int(num_items_lg2 + isz_lg2));
}

static void test_entry_possible_sizes(struct kunit *test)
{
        check_all_levels(test, test_lvl_possible_sizes, NULL);
}

static void sweep_all_pgsizes(struct kunit *test, struct pt_state *pts,
                              struct pt_write_attrs *attrs,
                              pt_oaddr_t test_oaddr)
{
        pt_vaddr_t pgsize_bitmap = pt_possible_sizes(pts);
        unsigned int isz_lg2 = pt_table_item_lg2sz(pts);
        unsigned int len_lg2;

        if (pts->index != 0)
                return;

        for (len_lg2 = 0; len_lg2 < PT_VADDR_MAX_LG2 - 1; len_lg2++) {
                struct pt_state sub_pts = *pts;
                pt_oaddr_t oaddr;

                if (!(pgsize_bitmap & log2_to_int(len_lg2)))
                        continue;

                oaddr = log2_set_mod(test_oaddr, 0, len_lg2);
                pt_install_leaf_entry(pts, oaddr, len_lg2, attrs);
                /* Verify that every contiguous item translates correctly */
                for (sub_pts.index = 0;
                     sub_pts.index != log2_to_int(len_lg2 - isz_lg2);
                     sub_pts.index++) {
                        KUNIT_ASSERT_PT_LOAD(test, &sub_pts, PT_ENTRY_OA);
                        KUNIT_ASSERT_EQ(test, pt_item_oa(&sub_pts),
                                        oaddr + sub_pts.index *
                                                        oalog2_mul(1, isz_lg2));
                        KUNIT_ASSERT_EQ(test, pt_entry_oa(&sub_pts), oaddr);
                        KUNIT_ASSERT_EQ(test, pt_entry_num_contig_lg2(&sub_pts),
                                        len_lg2 - isz_lg2);
                }

                pt_clear_entries(pts, len_lg2 - isz_lg2);
                KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_EMPTY);
        }
}

/*
 * Check that pt_install_leaf_entry() and pt_entry_oa() match.
 * Check that pt_clear_entries() works.
 */
static void test_lvl_entry_oa(struct kunit *test, struct pt_state *pts,
                              void *arg)
{
        unsigned int max_oa_lg2 = pts->range->common->max_oasz_lg2;
        struct kunit_iommu_priv *priv = test->priv;
        struct pt_write_attrs attrs = {};

        if (!pt_can_have_leaf(pts))
                return;

        KUNIT_ASSERT_NO_ERRNO_FN(test, "pt_iommu_set_prot",
                                 pt_iommu_set_prot(pts->range->common, &attrs,
                                                   IOMMU_READ));

        sweep_all_pgsizes(test, pts, &attrs, priv->test_oa);

        /* Check that the table can store the boundary OAs */
        sweep_all_pgsizes(test, pts, &attrs, 0);
        if (max_oa_lg2 == PT_OADDR_MAX_LG2)
                sweep_all_pgsizes(test, pts, &attrs, PT_OADDR_MAX);
        else
                sweep_all_pgsizes(test, pts, &attrs,
                                  oalog2_to_max_int(max_oa_lg2));
}

static void test_entry_oa(struct kunit *test)
{
        check_all_levels(test, test_lvl_entry_oa, NULL);
}

/* Test pt_attr_from_entry() */
static void test_lvl_attr_from_entry(struct kunit *test, struct pt_state *pts,
                                     void *arg)
{
        pt_vaddr_t pgsize_bitmap = pt_possible_sizes(pts);
        unsigned int isz_lg2 = pt_table_item_lg2sz(pts);
        struct kunit_iommu_priv *priv = test->priv;
        unsigned int len_lg2;
        unsigned int prot;

        if (!pt_can_have_leaf(pts))
                return;

        for (len_lg2 = 0; len_lg2 < PT_VADDR_MAX_LG2; len_lg2++) {
                if (!(pgsize_bitmap & log2_to_int(len_lg2)))
                        continue;
                for (prot = 0; prot <= (IOMMU_READ | IOMMU_WRITE | IOMMU_CACHE |
                                        IOMMU_NOEXEC | IOMMU_MMIO);
                     prot++) {
                        pt_oaddr_t oaddr;
                        struct pt_write_attrs attrs = {};
                        u64 good_entry;

                        /*
                         * If the format doesn't support this combination of
                         * prot bits skip it
                         */
                        if (pt_iommu_set_prot(pts->range->common, &attrs,
                                              prot)) {
                                /* But RW has to be supported */
                                KUNIT_ASSERT_NE(test, prot,
                                                IOMMU_READ | IOMMU_WRITE);
                                continue;
                        }

                        oaddr = log2_set_mod(priv->test_oa, 0, len_lg2);
                        pt_install_leaf_entry(pts, oaddr, len_lg2, &attrs);
                        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_OA);

                        good_entry = pts->entry;

                        memset(&attrs, 0, sizeof(attrs));
                        pt_attr_from_entry(pts, &attrs);

                        pt_clear_entries(pts, len_lg2 - isz_lg2);
                        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_EMPTY);

                        pt_install_leaf_entry(pts, oaddr, len_lg2, &attrs);
                        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_OA);

                        /*
                         * The descriptor produced by pt_attr_from_entry()
                         * produce an identical entry value when re-written
                         */
                        KUNIT_ASSERT_EQ(test, good_entry, pts->entry);

                        pt_clear_entries(pts, len_lg2 - isz_lg2);
                }
        }
}

static void test_attr_from_entry(struct kunit *test)
{
        check_all_levels(test, test_lvl_attr_from_entry, NULL);
}

static void test_lvl_dirty(struct kunit *test, struct pt_state *pts, void *arg)
{
        pt_vaddr_t pgsize_bitmap = pt_possible_sizes(pts);
        unsigned int isz_lg2 = pt_table_item_lg2sz(pts);
        struct kunit_iommu_priv *priv = test->priv;
        unsigned int start_idx = pts->index;
        struct pt_write_attrs attrs = {};
        unsigned int len_lg2;

        if (!pt_can_have_leaf(pts))
                return;

        KUNIT_ASSERT_NO_ERRNO_FN(test, "pt_iommu_set_prot",
                                 pt_iommu_set_prot(pts->range->common, &attrs,
                                                   IOMMU_READ | IOMMU_WRITE));

        for (len_lg2 = 0; len_lg2 < PT_VADDR_MAX_LG2; len_lg2++) {
                pt_oaddr_t oaddr;
                unsigned int i;

                if (!(pgsize_bitmap & log2_to_int(len_lg2)))
                        continue;

                oaddr = log2_set_mod(priv->test_oa, 0, len_lg2);
                pt_install_leaf_entry(pts, oaddr, len_lg2, &attrs);
                KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_OA);

                pt_load_entry(pts);
                pt_entry_make_write_clean(pts);
                pt_load_entry(pts);
                KUNIT_ASSERT_FALSE(test, pt_entry_is_write_dirty(pts));

                for (i = 0; i != log2_to_int(len_lg2 - isz_lg2); i++) {
                        /* dirty every contiguous entry */
                        pts->index = start_idx + i;
                        pt_load_entry(pts);
                        KUNIT_ASSERT_TRUE(test, pt_entry_make_write_dirty(pts));
                        pts->index = start_idx;
                        pt_load_entry(pts);
                        KUNIT_ASSERT_TRUE(test, pt_entry_is_write_dirty(pts));

                        pt_entry_make_write_clean(pts);
                        pt_load_entry(pts);
                        KUNIT_ASSERT_FALSE(test, pt_entry_is_write_dirty(pts));
                }

                pt_clear_entries(pts, len_lg2 - isz_lg2);
        }
}

static __maybe_unused void test_dirty(struct kunit *test)
{
        struct kunit_iommu_priv *priv = test->priv;

        if (!pt_dirty_supported(priv->common))
                kunit_skip(test,
                           "Page table features do not support dirty tracking");

        check_all_levels(test, test_lvl_dirty, NULL);
}

static void test_lvl_sw_bit_leaf(struct kunit *test, struct pt_state *pts,
                                 void *arg)
{
        struct kunit_iommu_priv *priv = test->priv;
        pt_vaddr_t pgsize_bitmap = pt_possible_sizes(pts);
        unsigned int isz_lg2 = pt_table_item_lg2sz(pts);
        struct pt_write_attrs attrs = {};
        unsigned int len_lg2;

        if (!pt_can_have_leaf(pts))
                return;
        if (pts->index != 0)
                return;

        KUNIT_ASSERT_NO_ERRNO_FN(test, "pt_iommu_set_prot",
                                 pt_iommu_set_prot(pts->range->common, &attrs,
                                                   IOMMU_READ));

        for (len_lg2 = 0; len_lg2 < PT_VADDR_MAX_LG2 - 1; len_lg2++) {
                pt_oaddr_t paddr = log2_set_mod(priv->test_oa, 0, len_lg2);
                struct pt_write_attrs new_attrs = {};
                unsigned int bitnr;

                if (!(pgsize_bitmap & log2_to_int(len_lg2)))
                        continue;

                pt_install_leaf_entry(pts, paddr, len_lg2, &attrs);

                for (bitnr = 0; bitnr <= pt_max_sw_bit(pts->range->common);
                     bitnr++)
                        KUNIT_ASSERT_FALSE(test,
                                           pt_test_sw_bit_acquire(pts, bitnr));

                for (bitnr = 0; bitnr <= pt_max_sw_bit(pts->range->common);
                     bitnr++) {
                        KUNIT_ASSERT_FALSE(test,
                                           pt_test_sw_bit_acquire(pts, bitnr));
                        pt_set_sw_bit_release(pts, bitnr);
                        KUNIT_ASSERT_TRUE(test,
                                          pt_test_sw_bit_acquire(pts, bitnr));
                }

                for (bitnr = 0; bitnr <= pt_max_sw_bit(pts->range->common);
                     bitnr++)
                        KUNIT_ASSERT_TRUE(test,
                                          pt_test_sw_bit_acquire(pts, bitnr));

                KUNIT_ASSERT_EQ(test, pt_item_oa(pts), paddr);

                /* SW bits didn't leak into the attrs */
                pt_attr_from_entry(pts, &new_attrs);
                KUNIT_ASSERT_MEMEQ(test, &new_attrs, &attrs, sizeof(attrs));

                pt_clear_entries(pts, len_lg2 - isz_lg2);
                KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_EMPTY);
        }
}

static __maybe_unused void test_sw_bit_leaf(struct kunit *test)
{
        check_all_levels(test, test_lvl_sw_bit_leaf, NULL);
}

static void test_lvl_sw_bit_table(struct kunit *test, struct pt_state *pts,
                                  void *arg)
{
        struct kunit_iommu_priv *priv = test->priv;
        struct pt_write_attrs attrs = {};
        pt_oaddr_t paddr =
                log2_set_mod(priv->test_oa, 0, priv->smallest_pgsz_lg2);
        unsigned int bitnr;

        if (!pt_can_have_leaf(pts))
                return;
        if (pts->index != 0)
                return;

        KUNIT_ASSERT_NO_ERRNO_FN(test, "pt_iommu_set_prot",
                                 pt_iommu_set_prot(pts->range->common, &attrs,
                                                   IOMMU_READ));

        KUNIT_ASSERT_TRUE(test, pt_install_table(pts, paddr, &attrs));

        for (bitnr = 0; bitnr <= pt_max_sw_bit(pts->range->common); bitnr++)
                KUNIT_ASSERT_FALSE(test, pt_test_sw_bit_acquire(pts, bitnr));

        for (bitnr = 0; bitnr <= pt_max_sw_bit(pts->range->common); bitnr++) {
                KUNIT_ASSERT_FALSE(test, pt_test_sw_bit_acquire(pts, bitnr));
                pt_set_sw_bit_release(pts, bitnr);
                KUNIT_ASSERT_TRUE(test, pt_test_sw_bit_acquire(pts, bitnr));
        }

        for (bitnr = 0; bitnr <= pt_max_sw_bit(pts->range->common); bitnr++)
                KUNIT_ASSERT_TRUE(test, pt_test_sw_bit_acquire(pts, bitnr));

        KUNIT_ASSERT_EQ(test, pt_table_pa(pts), paddr);

        pt_clear_entries(pts, ilog2(1));
        KUNIT_ASSERT_PT_LOAD(test, pts, PT_ENTRY_EMPTY);
}

static __maybe_unused void test_sw_bit_table(struct kunit *test)
{
        check_all_levels(test, test_lvl_sw_bit_table, NULL);
}

static struct kunit_case generic_pt_test_cases[] = {
        KUNIT_CASE_FMT(test_init),
        KUNIT_CASE_FMT(test_bitops),
        KUNIT_CASE_FMT(test_best_pgsize),
        KUNIT_CASE_FMT(test_table_ptr),
        KUNIT_CASE_FMT(test_max_va),
        KUNIT_CASE_FMT(test_table_radix),
        KUNIT_CASE_FMT(test_entry_possible_sizes),
        KUNIT_CASE_FMT(test_entry_oa),
        KUNIT_CASE_FMT(test_attr_from_entry),
#ifdef pt_entry_is_write_dirty
        KUNIT_CASE_FMT(test_dirty),
#endif
#ifdef pt_sw_bit
        KUNIT_CASE_FMT(test_sw_bit_leaf),
        KUNIT_CASE_FMT(test_sw_bit_table),
#endif
        {},
};

static int pt_kunit_generic_pt_init(struct kunit *test)
{
        struct kunit_iommu_priv *priv;
        int ret;

        priv = kunit_kzalloc(test, sizeof(*priv), GFP_KERNEL);
        if (!priv)
                return -ENOMEM;
        ret = pt_kunit_priv_init(test, priv);
        if (ret) {
                kunit_kfree(test, priv);
                return ret;
        }
        test->priv = priv;
        return 0;
}

static void pt_kunit_generic_pt_exit(struct kunit *test)
{
        struct kunit_iommu_priv *priv = test->priv;

        if (!test->priv)
                return;

        pt_iommu_deinit(priv->iommu);
        kunit_kfree(test, test->priv);
}

static struct kunit_suite NS(generic_pt_suite) = {
        .name = __stringify(NS(fmt_test)),
        .init = pt_kunit_generic_pt_init,
        .exit = pt_kunit_generic_pt_exit,
        .test_cases = generic_pt_test_cases,
};
kunit_test_suites(&NS(generic_pt_suite));