/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES
 */
#include "kunit_iommu.h"
#include "pt_iter.h"
#include <linux/generic_pt/iommu.h>
#include <linux/iommu.h>

static void do_map(struct kunit *test, pt_vaddr_t va, pt_oaddr_t pa,
		   pt_vaddr_t len);

struct count_valids {
	u64 per_size[PT_VADDR_MAX_LG2];
};

static int __count_valids(struct pt_range *range, void *arg, unsigned int level,
			  struct pt_table_p *table)
{
	struct pt_state pts = pt_init(range, level, table);
	struct count_valids *valids = arg;

	for_each_pt_level_entry(&pts) {
		if (pts.type == PT_ENTRY_TABLE) {
			pt_descend(&pts, arg, __count_valids);
			continue;
		}
		if (pts.type == PT_ENTRY_OA) {
			valids->per_size[pt_entry_oa_lg2sz(&pts)]++;
			continue;
		}
	}
	return 0;
}

/*
 * Number of valid table entries. This counts contiguous entries as a single
 * valid.
 */
static unsigned int count_valids(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range range = pt_top_range(priv->common);
	struct count_valids valids = {};
	u64 total = 0;
	unsigned int i;

	KUNIT_ASSERT_NO_ERRNO(test,
			      pt_walk_range(&range, __count_valids, &valids));

	for (i = 0; i != ARRAY_SIZE(valids.per_size); i++)
		total += valids.per_size[i];
	return total;
}

/* Only a single page size is present, count the number of valid entries */
static unsigned int count_valids_single(struct kunit *test, pt_vaddr_t pgsz)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range range = pt_top_range(priv->common);
	struct count_valids valids = {};
	u64 total = 0;
	unsigned int i;

	KUNIT_ASSERT_NO_ERRNO(test,
			      pt_walk_range(&range, __count_valids, &valids));

	for (i = 0; i != ARRAY_SIZE(valids.per_size); i++) {
		if ((1ULL << i) == pgsz)
			total = valids.per_size[i];
		else
			KUNIT_ASSERT_EQ(test, valids.per_size[i], 0);
	}
	return total;
}

static void do_unmap(struct kunit *test, pt_vaddr_t va, pt_vaddr_t len)
{
	struct kunit_iommu_priv *priv = test->priv;
	size_t ret;

	ret = iommu_unmap(&priv->domain, va, len);
	KUNIT_ASSERT_EQ(test, ret, len);
}

static void check_iova(struct kunit *test, pt_vaddr_t va, pt_oaddr_t pa,
		       pt_vaddr_t len)
{
	struct kunit_iommu_priv *priv = test->priv;
	pt_vaddr_t pfn = log2_div(va, priv->smallest_pgsz_lg2);
	pt_vaddr_t end_pfn = pfn + log2_div(len, priv->smallest_pgsz_lg2);

	for (; pfn != end_pfn; pfn++) {
		phys_addr_t res = iommu_iova_to_phys(&priv->domain,
						     pfn * priv->smallest_pgsz);

		KUNIT_ASSERT_EQ(test, res, (phys_addr_t)pa);
		if (res != pa)
			break;
		pa += priv->smallest_pgsz;
	}
}

static void test_increase_level(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_common *common = priv->common;

	if (!pt_feature(common, PT_FEAT_DYNAMIC_TOP))
		kunit_skip(test, "PT_FEAT_DYNAMIC_TOP not set for this format");

	if (IS_32BIT)
		kunit_skip(test, "Unable to test on 32bit");

	if (common->max_vasz_lg2 <= pt_top_range(common).max_vasz_lg2)
		kunit_skip(test,
			   "max_vasz_lg2 fits in starting level, no growth possible");

	/* Add every possible level to the max */
	while (common->max_vasz_lg2 != pt_top_range(common).max_vasz_lg2) {
		struct pt_range top_range = pt_top_range(common);

		if (top_range.va == 0)
			do_map(test, top_range.last_va + 1, 0,
			       priv->smallest_pgsz);
		else
			do_map(test, top_range.va - priv->smallest_pgsz, 0,
			       priv->smallest_pgsz);

		KUNIT_ASSERT_EQ(test, pt_top_range(common).top_level,
				top_range.top_level + 1);
		KUNIT_ASSERT_GE(test, common->max_vasz_lg2,
				pt_top_range(common).max_vasz_lg2);
	}
}

static void test_map_simple(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range range = pt_top_range(priv->common);
	struct count_valids valids = {};
	pt_vaddr_t pgsize_bitmap = priv->safe_pgsize_bitmap;
	unsigned int pgsz_lg2;
	pt_vaddr_t cur_va;

	/* Map every reported page size */
	cur_va = range.va + priv->smallest_pgsz * 256;
	for (pgsz_lg2 = 0; pgsz_lg2 != PT_VADDR_MAX_LG2; pgsz_lg2++) {
		pt_oaddr_t paddr = log2_set_mod(priv->test_oa, 0, pgsz_lg2);
		u64 len = log2_to_int(pgsz_lg2);

		if (!(pgsize_bitmap & len))
			continue;

		cur_va = ALIGN(cur_va, len);
		do_map(test, cur_va, paddr, len);
		if (len <= SZ_2G)
			check_iova(test, cur_va, paddr, len);
		cur_va += len;
	}

	/* The read interface reports that every page size was created */
	range = pt_top_range(priv->common);
	KUNIT_ASSERT_NO_ERRNO(test,
			      pt_walk_range(&range, __count_valids, &valids));
	for (pgsz_lg2 = 0; pgsz_lg2 != PT_VADDR_MAX_LG2; pgsz_lg2++) {
		if (pgsize_bitmap & (1ULL << pgsz_lg2))
			KUNIT_ASSERT_EQ(test, valids.per_size[pgsz_lg2], 1);
		else
			KUNIT_ASSERT_EQ(test, valids.per_size[pgsz_lg2], 0);
	}

	/* Unmap works */
	range = pt_top_range(priv->common);
	cur_va = range.va + priv->smallest_pgsz * 256;
	for (pgsz_lg2 = 0; pgsz_lg2 != PT_VADDR_MAX_LG2; pgsz_lg2++) {
		u64 len = log2_to_int(pgsz_lg2);

		if (!(pgsize_bitmap & len))
			continue;
		cur_va = ALIGN(cur_va, len);
		do_unmap(test, cur_va, len);
		cur_va += len;
	}
	KUNIT_ASSERT_EQ(test, count_valids(test), 0);
}

/*
 * Test to convert a table pointer into an OA by mapping something small,
 * unmapping it so as to leave behind a table pointer, then mapping something
 * larger that will convert the table into an OA.
 */
static void test_map_table_to_oa(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	pt_vaddr_t limited_pgbitmap =
		priv->info.pgsize_bitmap % (IS_32BIT ? SZ_2G : SZ_16G);
	struct pt_range range = pt_top_range(priv->common);
	unsigned int pgsz_lg2;
	pt_vaddr_t max_pgsize;
	pt_vaddr_t cur_va;

	max_pgsize = 1ULL << (vafls(limited_pgbitmap) - 1);
	KUNIT_ASSERT_TRUE(test, priv->info.pgsize_bitmap & max_pgsize);

	for (pgsz_lg2 = 0; pgsz_lg2 != PT_VADDR_MAX_LG2; pgsz_lg2++) {
		pt_oaddr_t paddr = log2_set_mod(priv->test_oa, 0, pgsz_lg2);
		u64 len = log2_to_int(pgsz_lg2);
		pt_vaddr_t offset;

		if (!(priv->info.pgsize_bitmap & len))
			continue;
		if (len > max_pgsize)
			break;

		cur_va = ALIGN(range.va + priv->smallest_pgsz * 256,
			       max_pgsize);
		for (offset = 0; offset != max_pgsize; offset += len)
			do_map(test, cur_va + offset, paddr + offset, len);
		check_iova(test, cur_va, paddr, max_pgsize);
		KUNIT_ASSERT_EQ(test, count_valids_single(test, len),
				log2_div(max_pgsize, pgsz_lg2));

		if (len == max_pgsize) {
			do_unmap(test, cur_va, max_pgsize);
		} else {
			do_unmap(test, cur_va, max_pgsize / 2);
			for (offset = max_pgsize / 2; offset != max_pgsize;
			     offset += len)
				do_unmap(test, cur_va + offset, len);
		}

		KUNIT_ASSERT_EQ(test, count_valids(test), 0);
	}
}

/*
 * Test unmapping a small page at the start of a large page. This always unmaps
 * the large page.
 */
static void test_unmap_split(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range top_range = pt_top_range(priv->common);
	pt_vaddr_t pgsize_bitmap = priv->safe_pgsize_bitmap;
	unsigned int pgsz_lg2;
	unsigned int count = 0;

	for (pgsz_lg2 = 0; pgsz_lg2 != PT_VADDR_MAX_LG2; pgsz_lg2++) {
		pt_vaddr_t base_len = log2_to_int(pgsz_lg2);
		unsigned int next_pgsz_lg2;

		if (!(pgsize_bitmap & base_len))
			continue;

		for (next_pgsz_lg2 = pgsz_lg2 + 1;
		     next_pgsz_lg2 != PT_VADDR_MAX_LG2; next_pgsz_lg2++) {
			pt_vaddr_t next_len = log2_to_int(next_pgsz_lg2);
			pt_vaddr_t vaddr = top_range.va;
			pt_oaddr_t paddr = 0;
			size_t gnmapped;

			if (!(pgsize_bitmap & next_len))
				continue;

			do_map(test, vaddr, paddr, next_len);
			gnmapped = iommu_unmap(&priv->domain, vaddr, base_len);
			KUNIT_ASSERT_EQ(test, gnmapped, next_len);

			/* Make sure unmap doesn't keep going */
			do_map(test, vaddr, paddr, next_len);
			do_map(test, vaddr + next_len, paddr, next_len);
			gnmapped = iommu_unmap(&priv->domain, vaddr, base_len);
			KUNIT_ASSERT_EQ(test, gnmapped, next_len);
			gnmapped = iommu_unmap(&priv->domain, vaddr + next_len,
					       next_len);
			KUNIT_ASSERT_EQ(test, gnmapped, next_len);

			count++;
		}
	}

	if (count == 0)
		kunit_skip(test, "Test needs two page sizes");
}

static void unmap_collisions(struct kunit *test, struct maple_tree *mt,
			     pt_vaddr_t start, pt_vaddr_t last)
{
	struct kunit_iommu_priv *priv = test->priv;
	MA_STATE(mas, mt, start, last);
	void *entry;

	mtree_lock(mt);
	mas_for_each(&mas, entry, last) {
		pt_vaddr_t mas_start = mas.index;
		pt_vaddr_t len = (mas.last - mas_start) + 1;
		pt_oaddr_t paddr;

		mas_erase(&mas);
		mas_pause(&mas);
		mtree_unlock(mt);

		paddr = oalog2_mod(mas_start, priv->common->max_oasz_lg2);
		check_iova(test, mas_start, paddr, len);
		do_unmap(test, mas_start, len);
		mtree_lock(mt);
	}
	mtree_unlock(mt);
}

static void clamp_range(struct kunit *test, struct pt_range *range)
{
	struct kunit_iommu_priv *priv = test->priv;

	if (range->last_va - range->va > SZ_1G)
		range->last_va = range->va + SZ_1G;
	KUNIT_ASSERT_NE(test, range->last_va, PT_VADDR_MAX);
	if (range->va <= MAPLE_RESERVED_RANGE)
		range->va =
			ALIGN(MAPLE_RESERVED_RANGE, priv->smallest_pgsz);
}

/*
 * Randomly map and unmap ranges that can large physical pages. If a random
 * range overlaps with existing ranges then unmap them. This hits all the
 * special cases.
 */
static void test_random_map(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range upper_range = pt_upper_range(priv->common);
	struct pt_range top_range = pt_top_range(priv->common);
	struct maple_tree mt;
	unsigned int iter;

	mt_init(&mt);

	/*
	 * Shrink the range so randomization is more likely to have
	 * intersections
	 */
	clamp_range(test, &top_range);
	clamp_range(test, &upper_range);

	for (iter = 0; iter != 1000; iter++) {
		struct pt_range *range = &top_range;
		pt_oaddr_t paddr;
		pt_vaddr_t start;
		pt_vaddr_t end;
		int ret;

		if (pt_feature(priv->common, PT_FEAT_SIGN_EXTEND) &&
		    ULONG_MAX >= PT_VADDR_MAX && get_random_u32_inclusive(0, 1))
			range = &upper_range;

		start = get_random_u32_below(
			min(U32_MAX, range->last_va - range->va));
		end = get_random_u32_below(
			min(U32_MAX, range->last_va - start));

		start = ALIGN_DOWN(start, priv->smallest_pgsz);
		end = ALIGN(end, priv->smallest_pgsz);
		start += range->va;
		end += start;
		if (start < range->va || end > range->last_va + 1 ||
		    start >= end)
			continue;

		/* Try overmapping to test the failure handling */
		paddr = oalog2_mod(start, priv->common->max_oasz_lg2);
		ret = iommu_map(&priv->domain, start, paddr, end - start,
				IOMMU_READ | IOMMU_WRITE, GFP_KERNEL);
		if (ret) {
			KUNIT_ASSERT_EQ(test, ret, -EADDRINUSE);
			unmap_collisions(test, &mt, start, end - 1);
			do_map(test, start, paddr, end - start);
		}

		KUNIT_ASSERT_NO_ERRNO_FN(test, "mtree_insert_range",
					 mtree_insert_range(&mt, start, end - 1,
							    XA_ZERO_ENTRY,
							    GFP_KERNEL));

		check_iova(test, start, paddr, end - start);
		if (iter % 100)
			cond_resched();
	}

	unmap_collisions(test, &mt, 0, PT_VADDR_MAX);
	KUNIT_ASSERT_EQ(test, count_valids(test), 0);

	mtree_destroy(&mt);
}

/* See https://lore.kernel.org/r/b9b18a03-63a2-4065-a27e-d92dd5c860bc@amd.com */
static void test_pgsize_boundary(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range top_range = pt_top_range(priv->common);

	if (top_range.va != 0 || top_range.last_va < 0xfef9ffff ||
	    priv->smallest_pgsz != SZ_4K)
		kunit_skip(test, "Format does not have the required range");

	do_map(test, 0xfef80000, 0x208b95d000, 0xfef9ffff - 0xfef80000 + 1);
}

/* See https://lore.kernel.org/r/20250826143816.38686-1-eugkoira@amazon.com */
static void test_mixed(struct kunit *test)
{
	struct kunit_iommu_priv *priv = test->priv;
	struct pt_range top_range = pt_top_range(priv->common);
	u64 start = 0x3fe400ULL << 12;
	u64 end = 0x4c0600ULL << 12;
	pt_vaddr_t len = end - start;
	pt_oaddr_t oa = start;

	if (top_range.last_va <= start || sizeof(unsigned long) == 4)
		kunit_skip(test, "range is too small");
	if ((priv->safe_pgsize_bitmap & GENMASK(30, 21)) != (BIT(30) | BIT(21)))
		kunit_skip(test, "incompatible psize");

	do_map(test, start, oa, len);
	/* 14 2M, 3 1G, 3 2M */
	KUNIT_ASSERT_EQ(test, count_valids(test), 20);
	check_iova(test, start, oa, len);
}

static struct kunit_case iommu_test_cases[] = {
	KUNIT_CASE_FMT(test_increase_level),
	KUNIT_CASE_FMT(test_map_simple),
	KUNIT_CASE_FMT(test_map_table_to_oa),
	KUNIT_CASE_FMT(test_unmap_split),
	KUNIT_CASE_FMT(test_random_map),
	KUNIT_CASE_FMT(test_pgsize_boundary),
	KUNIT_CASE_FMT(test_mixed),
	{},
};

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

	priv = kunit_kzalloc(test, sizeof(*priv), GFP_KERNEL);
	if (!priv)
		return -ENOMEM;

	priv->orig_nr_secondary_pagetable =
		global_node_page_state(NR_SECONDARY_PAGETABLE);
	ret = pt_kunit_priv_init(test, priv);
	if (ret) {
		kunit_kfree(test, priv);
		return ret;
	}
	test->priv = priv;
	return 0;
}

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

	if (!test->priv)
		return;

	pt_iommu_deinit(priv->iommu);
	/*
	 * Look for memory leaks, assumes kunit is running isolated and nothing
	 * else is using secondary page tables.
	 */
	KUNIT_ASSERT_EQ(test, priv->orig_nr_secondary_pagetable,
			global_node_page_state(NR_SECONDARY_PAGETABLE));
	kunit_kfree(test, test->priv);
}

static struct kunit_suite NS(iommu_suite) = {
	.name = __stringify(NS(iommu_test)),
	.init = pt_kunit_iommu_init,
	.exit = pt_kunit_iommu_exit,
	.test_cases = iommu_test_cases,
};
kunit_test_suites(&NS(iommu_suite));

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Kunit for generic page table");
MODULE_IMPORT_NS("GENERIC_PT_IOMMU");