/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES
 */
#ifndef __GENERIC_PT_KUNIT_IOMMU_H
#define __GENERIC_PT_KUNIT_IOMMU_H

#define GENERIC_PT_KUNIT 1
#include <kunit/device.h>
#include <kunit/test.h>
#include "../iommu-pages.h"
#include "pt_iter.h"

#define pt_iommu_table_cfg CONCATENATE(pt_iommu_table, _cfg)
#define pt_iommu_init CONCATENATE(CONCATENATE(pt_iommu_, PTPFX), init)
int pt_iommu_init(struct pt_iommu_table *fmt_table,
		  const struct pt_iommu_table_cfg *cfg, gfp_t gfp);

/* The format can provide a list of configurations it would like to test */
#ifdef kunit_fmt_cfgs
static const void *kunit_pt_gen_params_cfg(struct kunit *test, const void *prev,
					   char *desc)
{
	uintptr_t cfg_id = (uintptr_t)prev;

	cfg_id++;
	if (cfg_id >= ARRAY_SIZE(kunit_fmt_cfgs) + 1)
		return NULL;
	snprintf(desc, KUNIT_PARAM_DESC_SIZE, "%s_cfg_%u",
		 __stringify(PTPFX_RAW), (unsigned int)(cfg_id - 1));
	return (void *)cfg_id;
}
#define KUNIT_CASE_FMT(test_name) \
	KUNIT_CASE_PARAM(test_name, kunit_pt_gen_params_cfg)
#else
#define KUNIT_CASE_FMT(test_name) KUNIT_CASE(test_name)
#endif

#define KUNIT_ASSERT_NO_ERRNO(test, ret)                                       \
	KUNIT_ASSERT_EQ_MSG(test, ret, 0, KUNIT_SUBSUBTEST_INDENT "errno %pe", \
			    ERR_PTR(ret))

#define KUNIT_ASSERT_NO_ERRNO_FN(test, fn, ret)                          \
	KUNIT_ASSERT_EQ_MSG(test, ret, 0,                                \
			    KUNIT_SUBSUBTEST_INDENT "errno %pe from %s", \
			    ERR_PTR(ret), fn)

/*
 * When the test is run on a 32 bit system unsigned long can be 32 bits. This
 * cause the iommu op signatures to be restricted to 32 bits. Meaning the test
 * has to be mindful not to create any VA's over the 32 bit limit. Reduce the
 * scope of the testing as the main purpose of checking on full 32 bit is to
 * look for 32bitism in the core code. Run the test on i386 with X86_PAE=y to
 * get the full coverage when dma_addr_t & phys_addr_t are 8 bytes
 */
#define IS_32BIT (sizeof(unsigned long) == 4)

struct kunit_iommu_priv {
	union {
		struct iommu_domain domain;
		struct pt_iommu_table fmt_table;
	};
	spinlock_t top_lock;
	struct device *dummy_dev;
	struct pt_iommu *iommu;
	struct pt_common *common;
	struct pt_iommu_table_cfg cfg;
	struct pt_iommu_info info;
	unsigned int smallest_pgsz_lg2;
	pt_vaddr_t smallest_pgsz;
	unsigned int largest_pgsz_lg2;
	pt_oaddr_t test_oa;
	pt_vaddr_t safe_pgsize_bitmap;
	unsigned long orig_nr_secondary_pagetable;

};
PT_IOMMU_CHECK_DOMAIN(struct kunit_iommu_priv, fmt_table.iommu, domain);

static void pt_kunit_iotlb_sync(struct iommu_domain *domain,
				struct iommu_iotlb_gather *gather)
{
	iommu_put_pages_list(&gather->freelist);
}

#define IOMMU_PT_DOMAIN_OPS1(x) IOMMU_PT_DOMAIN_OPS(x)
static const struct iommu_domain_ops kunit_pt_ops = {
	IOMMU_PT_DOMAIN_OPS1(PTPFX_RAW),
	.iotlb_sync = &pt_kunit_iotlb_sync,
};

static void pt_kunit_change_top(struct pt_iommu *iommu_table,
				phys_addr_t top_paddr, unsigned int top_level)
{
}

static spinlock_t *pt_kunit_get_top_lock(struct pt_iommu *iommu_table)
{
	struct kunit_iommu_priv *priv = container_of(
		iommu_table, struct kunit_iommu_priv, fmt_table.iommu);

	return &priv->top_lock;
}

static const struct pt_iommu_driver_ops pt_kunit_driver_ops = {
	.change_top = &pt_kunit_change_top,
	.get_top_lock = &pt_kunit_get_top_lock,
};

static int pt_kunit_priv_init(struct kunit *test, struct kunit_iommu_priv *priv)
{
	unsigned int va_lg2sz;
	int ret;

	/* Enough so the memory allocator works */
	priv->dummy_dev = kunit_device_register(test, "pt_kunit_dev");
	if (IS_ERR(priv->dummy_dev))
		return PTR_ERR(priv->dummy_dev);
	set_dev_node(priv->dummy_dev, NUMA_NO_NODE);

	spin_lock_init(&priv->top_lock);

#ifdef kunit_fmt_cfgs
	priv->cfg = kunit_fmt_cfgs[((uintptr_t)test->param_value) - 1];
	/*
	 * The format can set a list of features that the kunit_fmt_cfgs
	 * controls, other features are default to on.
	 */
	priv->cfg.common.features |= PT_SUPPORTED_FEATURES &
				     (~KUNIT_FMT_FEATURES);
#else
	priv->cfg.common.features = PT_SUPPORTED_FEATURES;
#endif

	/* Defaults, for the kunit */
	if (!priv->cfg.common.hw_max_vasz_lg2)
		priv->cfg.common.hw_max_vasz_lg2 = PT_MAX_VA_ADDRESS_LG2;
	if (!priv->cfg.common.hw_max_oasz_lg2)
		priv->cfg.common.hw_max_oasz_lg2 = pt_max_oa_lg2(NULL);

	priv->fmt_table.iommu.nid = NUMA_NO_NODE;
	priv->fmt_table.iommu.driver_ops = &pt_kunit_driver_ops;
	priv->fmt_table.iommu.iommu_device = priv->dummy_dev;
	priv->domain.ops = &kunit_pt_ops;
	ret = pt_iommu_init(&priv->fmt_table, &priv->cfg, GFP_KERNEL);
	if (ret) {
		if (ret == -EOVERFLOW)
			kunit_skip(
				test,
				"This configuration cannot be tested on 32 bit");
		return ret;
	}

	priv->iommu = &priv->fmt_table.iommu;
	priv->common = common_from_iommu(&priv->fmt_table.iommu);
	priv->iommu->ops->get_info(priv->iommu, &priv->info);

	/*
	 * size_t is used to pass the mapping length, it can be 32 bit, truncate
	 * the pagesizes so we don't use large sizes.
	 */
	priv->info.pgsize_bitmap = (size_t)priv->info.pgsize_bitmap;

	priv->smallest_pgsz_lg2 = vaffs(priv->info.pgsize_bitmap);
	priv->smallest_pgsz = log2_to_int(priv->smallest_pgsz_lg2);
	priv->largest_pgsz_lg2 =
		vafls((dma_addr_t)priv->info.pgsize_bitmap) - 1;

	priv->test_oa =
		oalog2_mod(0x74a71445deadbeef, priv->common->max_oasz_lg2);

	/*
	 * We run out of VA space if the mappings get too big, make something
	 * smaller that can safely pass through dma_addr_t API.
	 */
	va_lg2sz = priv->common->max_vasz_lg2;
	if (IS_32BIT && va_lg2sz > 32)
		va_lg2sz = 32;
	priv->safe_pgsize_bitmap =
		log2_mod(priv->info.pgsize_bitmap, va_lg2sz - 1);

	return 0;
}

#endif