/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES
 *
 * "Templated C code" for implementing the iommu operations for page tables.
 * This is compiled multiple times, over all the page table formats to pick up
 * the per-format definitions.
 */
#ifndef __GENERIC_PT_IOMMU_PT_H
#define __GENERIC_PT_IOMMU_PT_H

#include "pt_iter.h"

#include <linux/export.h>
#include <linux/iommu.h>
#include "../iommu-pages.h"
#include <linux/cleanup.h>
#include <linux/dma-mapping.h>

enum {
	SW_BIT_CACHE_FLUSH_DONE = 0,
};

static void flush_writes_range(const struct pt_state *pts,
			       unsigned int start_index, unsigned int end_index)
{
	if (pts_feature(pts, PT_FEAT_DMA_INCOHERENT))
		iommu_pages_flush_incoherent(
			iommu_from_common(pts->range->common)->iommu_device,
			pts->table, start_index * PT_ITEM_WORD_SIZE,
			(end_index - start_index) * PT_ITEM_WORD_SIZE);
}

static void flush_writes_item(const struct pt_state *pts)
{
	if (pts_feature(pts, PT_FEAT_DMA_INCOHERENT))
		iommu_pages_flush_incoherent(
			iommu_from_common(pts->range->common)->iommu_device,
			pts->table, pts->index * PT_ITEM_WORD_SIZE,
			PT_ITEM_WORD_SIZE);
}

struct iommupt_pending_gather {
	struct iommu_iotlb_gather *iotlb_gather;
	struct iommu_pages_list free_list;
	u8 leaf_levels_bitmap;
	u8 table_levels_bitmap;
};

static void gather_add_table(struct iommupt_pending_gather *pending,
			     const struct pt_state *pts,
			     struct pt_table_p *table)
{
	iommu_pages_list_add(&pending->free_list, table);
	if (pts_feature(pts, PT_FEAT_DETAILED_GATHER))
		pending->table_levels_bitmap |= BIT(pts->level);
}

static void gather_add_leaf(struct iommupt_pending_gather *pending,
			    const struct pt_state *pts)
{
	if (!pts_feature(pts, PT_FEAT_DETAILED_GATHER))
		return;

	pending->leaf_levels_bitmap |= BIT(pts->level);
}

static void gather_range_pending(struct iommupt_pending_gather *pending,
				 struct pt_iommu *iommu_table, pt_vaddr_t iova,
				 pt_vaddr_t len)
{
	struct iommu_iotlb_gather *iotlb_gather = pending->iotlb_gather;
	struct pt_common *common = common_from_iommu(iommu_table);

	if (pt_feature(common, PT_FEAT_DMA_INCOHERENT))
		iommu_pages_stop_incoherent_list(&pending->free_list,
						 iommu_table->iommu_device);

	/*
	 * If running in DMA-FQ mode then the unmap will be followed by an IOTLB
	 * flush all so we need to optimize by never flushing the IOTLB here.
	 *
	 * For NO_GAPS the user gets to pick if flushing all or doing micro
	 * flushes is better for their work load by choosing DMA vs DMA-FQ
	 * operation. Drivers should also see shadow_on_flush.
	 */
	if (!iommu_iotlb_gather_queued(iotlb_gather)) {
		if (pt_feature(common, PT_FEAT_FLUSH_RANGE_NO_GAPS) &&
		    iommu_iotlb_gather_is_disjoint(iotlb_gather, iova, len)) {
			iommu_iotlb_sync(&iommu_table->domain, iotlb_gather);
			/*
			 * Note that the sync frees the gather's free list, so
			 * we must not have any pages on that list that are
			 * covered by iova/len
			 */
		}
		iommu_iotlb_gather_add_range(iotlb_gather, iova, len);
	}

	iommu_pages_list_splice(&pending->free_list, &iotlb_gather->freelist);
	INIT_LIST_HEAD(&pending->free_list.pages);

	if (pt_feature(common, PT_FEAT_DETAILED_GATHER)) {
		iotlb_gather->pt.leaf_levels_bitmap |=
			pending->leaf_levels_bitmap;
		iotlb_gather->pt.table_levels_bitmap |=
			pending->table_levels_bitmap;
		pending->leaf_levels_bitmap = 0;
		pending->table_levels_bitmap = 0;
	}
}

#define DOMAIN_NS(op) CONCATENATE(CONCATENATE(pt_iommu_, PTPFX), op)

static int make_range_ul(struct pt_common *common, struct pt_range *range,
			 unsigned long iova, unsigned long len)
{
	unsigned long last;

	if (unlikely(len == 0))
		return -EINVAL;

	if (check_add_overflow(iova, len - 1, &last))
		return -EOVERFLOW;

	*range = pt_make_range(common, iova, last);
	if (sizeof(iova) > sizeof(range->va)) {
		if (unlikely(range->va != iova || range->last_va != last))
			return -EOVERFLOW;
	}
	return 0;
}

static __maybe_unused int make_range_u64(struct pt_common *common,
					 struct pt_range *range, u64 iova,
					 u64 len)
{
	if (unlikely(iova > ULONG_MAX || len > ULONG_MAX))
		return -EOVERFLOW;
	return make_range_ul(common, range, iova, len);
}

/*
 * Some APIs use unsigned long, while othersuse dma_addr_t as the type. Dispatch
 * to the correct validation based on the type.
 */
#define make_range_no_check(common, range, iova, len)                   \
	({                                                              \
		int ret;                                                \
		if (sizeof(iova) > sizeof(unsigned long) ||             \
		    sizeof(len) > sizeof(unsigned long))                \
			ret = make_range_u64(common, range, iova, len); \
		else                                                    \
			ret = make_range_ul(common, range, iova, len);  \
		ret;                                                    \
	})

#define make_range(common, range, iova, len)                             \
	({                                                               \
		int ret = make_range_no_check(common, range, iova, len); \
		if (!ret)                                                \
			ret = pt_check_range(range);                     \
		ret;                                                     \
	})

static inline unsigned int compute_best_pgsize(struct pt_state *pts,
					       pt_oaddr_t oa)
{
	struct pt_iommu *iommu_table = iommu_from_common(pts->range->common);

	if (!pt_can_have_leaf(pts))
		return 0;

	/*
	 * The page size is limited by the domain's bitmap. This allows the core
	 * code to reduce the supported page sizes by changing the bitmap.
	 */
	return pt_compute_best_pgsize(pt_possible_sizes(pts) &
					      iommu_table->domain.pgsize_bitmap,
				      pts->range->va, pts->range->last_va, oa);
}

static __always_inline int __do_iova_to_phys(struct pt_range *range, void *arg,
					     unsigned int level,
					     struct pt_table_p *table,
					     pt_level_fn_t descend_fn)
{
	struct pt_state pts = pt_init(range, level, table);
	pt_oaddr_t *res = arg;

	switch (pt_load_single_entry(&pts)) {
	case PT_ENTRY_EMPTY:
		return -ENOENT;
	case PT_ENTRY_TABLE:
		return pt_descend(&pts, arg, descend_fn);
	case PT_ENTRY_OA:
		*res = pt_entry_oa_exact(&pts);
		return 0;
	}
	return -ENOENT;
}
PT_MAKE_LEVELS(__iova_to_phys, __do_iova_to_phys);

/**
 * iova_to_phys() - Return the output address for the given IOVA
 * @domain: Table to query
 * @iova: IO virtual address to query
 *
 * Determine the output address from the given IOVA. @iova may have any
 * alignment, the returned physical will be adjusted with any sub page offset.
 *
 * Context: The caller must hold a read range lock that includes @iova.
 *
 * Return: 0 if there is no translation for the given iova.
 */
phys_addr_t DOMAIN_NS(iova_to_phys)(struct iommu_domain *domain,
				    dma_addr_t iova)
{
	struct pt_iommu *iommu_table =
		container_of(domain, struct pt_iommu, domain);
	struct pt_range range;
	pt_oaddr_t res;
	int ret;

	ret = make_range(common_from_iommu(iommu_table), &range, iova, 1);
	if (ret)
		return 0;

	ret = pt_walk_range(&range, __iova_to_phys, &res);
	/* PHYS_ADDR_MAX would be a better error code */
	if (ret)
		return 0;
	return res;
}
EXPORT_SYMBOL_NS_GPL(DOMAIN_NS(iova_to_phys), "GENERIC_PT_IOMMU");

struct pt_iommu_dirty_args {
	struct iommu_dirty_bitmap *dirty;
	unsigned int flags;
};

static void record_dirty(struct pt_state *pts,
			 struct pt_iommu_dirty_args *dirty,
			 unsigned int num_contig_lg2)
{
	pt_vaddr_t dirty_len;

	if (num_contig_lg2 != ilog2(1)) {
		unsigned int index = pts->index;
		unsigned int end_index = log2_set_mod_max_t(
			unsigned int, pts->index, num_contig_lg2);

		/* Adjust for being contained inside a contiguous page */
		end_index = min(end_index, pts->end_index);
		dirty_len = (end_index - index) *
				log2_to_int(pt_table_item_lg2sz(pts));
	} else {
		dirty_len = log2_to_int(pt_table_item_lg2sz(pts));
	}

	if (dirty->dirty->bitmap)
		iova_bitmap_set(dirty->dirty->bitmap, pts->range->va,
				dirty_len);

	if (!(dirty->flags & IOMMU_DIRTY_NO_CLEAR)) {
		/*
		 * No write log required because DMA incoherence and atomic
		 * dirty tracking bits can't work together
		 */
		pt_entry_make_write_clean(pts);
		iommu_iotlb_gather_add_range(dirty->dirty->gather,
					     pts->range->va, dirty_len);
	}
}

static inline int __read_and_clear_dirty(struct pt_range *range, void *arg,
					 unsigned int level,
					 struct pt_table_p *table)
{
	struct pt_state pts = pt_init(range, level, table);
	struct pt_iommu_dirty_args *dirty = arg;
	int ret;

	for_each_pt_level_entry(&pts) {
		if (pts.type == PT_ENTRY_TABLE) {
			ret = pt_descend(&pts, arg, __read_and_clear_dirty);
			if (ret)
				return ret;
			continue;
		}
		if (pts.type == PT_ENTRY_OA && pt_entry_is_write_dirty(&pts))
			record_dirty(&pts, dirty,
				     pt_entry_num_contig_lg2(&pts));
	}
	return 0;
}

/**
 * read_and_clear_dirty() - Manipulate the HW set write dirty state
 * @domain: Domain to manipulate
 * @iova: IO virtual address to start
 * @size: Length of the IOVA
 * @flags: A bitmap of IOMMU_DIRTY_NO_CLEAR
 * @dirty: Place to store the dirty bits
 *
 * Iterate over all the entries in the mapped range and record their write dirty
 * status in iommu_dirty_bitmap. If IOMMU_DIRTY_NO_CLEAR is not specified then
 * the entries will be left dirty, otherwise they are returned to being not
 * write dirty.
 *
 * Context: The caller must hold a read range lock that includes @iova.
 *
 * Returns: -ERRNO on failure, 0 on success.
 */
int DOMAIN_NS(read_and_clear_dirty)(struct iommu_domain *domain,
				    unsigned long iova, size_t size,
				    unsigned long flags,
				    struct iommu_dirty_bitmap *dirty)
{
	struct pt_iommu *iommu_table =
		container_of(domain, struct pt_iommu, domain);
	struct pt_iommu_dirty_args dirty_args = {
		.dirty = dirty,
		.flags = flags,
	};
	struct pt_range range;
	int ret;

#if !IS_ENABLED(CONFIG_IOMMUFD_DRIVER) || !defined(pt_entry_is_write_dirty)
	return -EOPNOTSUPP;
#endif

	ret = make_range(common_from_iommu(iommu_table), &range, iova, size);
	if (ret)
		return ret;

	ret = pt_walk_range(&range, __read_and_clear_dirty, &dirty_args);
	PT_WARN_ON(ret);
	return ret;
}
EXPORT_SYMBOL_NS_GPL(DOMAIN_NS(read_and_clear_dirty), "GENERIC_PT_IOMMU");

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

	switch (pt_load_single_entry(&pts)) {
	case PT_ENTRY_EMPTY:
		return -ENOENT;
	case PT_ENTRY_TABLE:
		return pt_descend(&pts, arg, __set_dirty);
	case PT_ENTRY_OA:
		if (!pt_entry_make_write_dirty(&pts))
			return -EAGAIN;
		return 0;
	}
	return -ENOENT;
}

static int __maybe_unused NS(set_dirty)(struct pt_iommu *iommu_table,
					dma_addr_t iova)
{
	struct pt_range range;
	int ret;

	ret = make_range(common_from_iommu(iommu_table), &range, iova, 1);
	if (ret)
		return ret;

	/*
	 * Note: There is no locking here yet, if the test suite races this it
	 * can crash. It should use RCU locking eventually.
	 */
	return pt_walk_range(&range, __set_dirty, NULL);
}

struct pt_iommu_collect_args {
	struct iommupt_pending_gather pending;
	/* Fail if any OAs are within the range */
	u8 check_mapped : 1;
};

static int __collect_tables(struct pt_range *range, void *arg,
			    unsigned int level, struct pt_table_p *table)
{
	struct pt_state pts = pt_init(range, level, table);
	struct pt_iommu_collect_args *collect = arg;
	int ret;

	if (!collect->check_mapped && !pt_can_have_table(&pts))
		return 0;

	for_each_pt_level_entry(&pts) {
		if (pts.type == PT_ENTRY_TABLE) {
			gather_add_table(&collect->pending, &pts,
					 pts.table_lower);
			ret = pt_descend(&pts, arg, __collect_tables);
			if (ret)
				return ret;
			continue;
		}
		if (pts.type == PT_ENTRY_OA && collect->check_mapped)
			return -EADDRINUSE;
	}
	return 0;
}

enum alloc_mode {ALLOC_NORMAL, ALLOC_DEFER_COHERENT_FLUSH};

/* Allocate a table, the empty table will be ready to be installed. */
static inline struct pt_table_p *_table_alloc(struct pt_common *common,
					      size_t lg2sz, gfp_t gfp,
					      enum alloc_mode mode)
{
	struct pt_iommu *iommu_table = iommu_from_common(common);
	struct pt_table_p *table_mem;

	table_mem = iommu_alloc_pages_node_sz(iommu_table->nid, gfp,
					      log2_to_int(lg2sz));
	if (!table_mem)
		return ERR_PTR(-ENOMEM);

	if (pt_feature(common, PT_FEAT_DMA_INCOHERENT) &&
	    mode == ALLOC_NORMAL) {
		int ret = iommu_pages_start_incoherent(
			table_mem, iommu_table->iommu_device);
		if (ret) {
			iommu_free_pages(table_mem);
			return ERR_PTR(ret);
		}
	}
	return table_mem;
}

static inline struct pt_table_p *table_alloc_top(struct pt_common *common,
						 uintptr_t top_of_table,
						 gfp_t gfp,
						 enum alloc_mode mode)
{
	/*
	 * Top doesn't need the free list or otherwise, so it technically
	 * doesn't need to use iommu pages. Use the API anyhow as the top is
	 * usually not smaller than PAGE_SIZE to keep things simple.
	 */
	return _table_alloc(common, pt_top_memsize_lg2(common, top_of_table),
			    gfp, mode);
}

/* Allocate an interior table */
static inline struct pt_table_p *table_alloc(const struct pt_state *parent_pts,
					     gfp_t gfp, enum alloc_mode mode)
{
	struct pt_state child_pts =
		pt_init(parent_pts->range, parent_pts->level - 1, NULL);

	return _table_alloc(parent_pts->range->common,
			    pt_num_items_lg2(&child_pts) +
				    ilog2(PT_ITEM_WORD_SIZE),
			    gfp, mode);
}

static inline int pt_iommu_new_table(struct pt_state *pts,
				     struct pt_write_attrs *attrs)
{
	struct pt_table_p *table_mem;
	phys_addr_t phys;

	/* Given PA/VA/length can't be represented */
	if (PT_WARN_ON(!pt_can_have_table(pts)))
		return -ENXIO;

	table_mem = table_alloc(pts, attrs->gfp, ALLOC_NORMAL);
	if (IS_ERR(table_mem))
		return PTR_ERR(table_mem);

	phys = virt_to_phys(table_mem);
	if (!pt_install_table(pts, phys, attrs)) {
		iommu_pages_free_incoherent(
			table_mem,
			iommu_from_common(pts->range->common)->iommu_device);
		return -EAGAIN;
	}

	if (pts_feature(pts, PT_FEAT_DMA_INCOHERENT)) {
		flush_writes_item(pts);
		pt_set_sw_bit_release(pts, SW_BIT_CACHE_FLUSH_DONE);
	}

	if (IS_ENABLED(CONFIG_DEBUG_GENERIC_PT)) {
		/*
		 * The underlying table can't store the physical table address.
		 * This happens when kunit testing tables outside their normal
		 * environment where a CPU might be limited.
		 */
		pt_load_single_entry(pts);
		if (PT_WARN_ON(pt_table_pa(pts) != phys)) {
			pt_clear_entries(pts, ilog2(1));
			iommu_pages_free_incoherent(
				table_mem, iommu_from_common(pts->range->common)
						   ->iommu_device);
			return -EINVAL;
		}
	}

	pts->table_lower = table_mem;
	return 0;
}

struct pt_iommu_map_args {
	struct iommu_iotlb_gather *iotlb_gather;
	struct pt_write_attrs attrs;
	pt_oaddr_t oa;
	unsigned int leaf_pgsize_lg2;
	unsigned int leaf_level;
	pt_vaddr_t num_leaves;
};

/*
 * This will recursively check any tables in the block to validate they are
 * empty and then free them through the gather.
 */
static int clear_contig(const struct pt_state *start_pts,
			struct iommu_iotlb_gather *iotlb_gather,
			unsigned int step, unsigned int pgsize_lg2)
{
	struct pt_iommu *iommu_table =
		iommu_from_common(start_pts->range->common);
	struct pt_range range = *start_pts->range;
	struct pt_state pts =
		pt_init(&range, start_pts->level, start_pts->table);
	struct pt_iommu_collect_args collect = {
		.check_mapped = true,
		.pending.iotlb_gather = iotlb_gather,
		.pending.free_list = IOMMU_PAGES_LIST_INIT(
			collect.pending.free_list),
	};
	int ret;

	pts.index = start_pts->index;
	pts.end_index = start_pts->index + step;
	for (; _pt_iter_load(&pts); pt_next_entry(&pts)) {
		if (pts.type == PT_ENTRY_TABLE) {
			ret = pt_walk_descend_all(&pts, __collect_tables,
						  &collect);
			if (ret)
				return ret;

			/*
			 * The table item must be cleared before we can update
			 * the gather
			 */
			pt_clear_entries(&pts, ilog2(1));
			flush_writes_item(&pts);

			gather_add_table(&collect.pending, &pts,
					 pts.table_lower);
			gather_range_pending(
				&collect.pending, iommu_table, range.va,
				log2_to_int(pt_table_item_lg2sz(&pts)));
		} else if (pts.type != PT_ENTRY_EMPTY) {
			return -EADDRINUSE;
		}
	}
	return 0;
}

static int __map_range_leaf(struct pt_range *range, void *arg,
			    unsigned int level, struct pt_table_p *table)
{
	struct pt_iommu *iommu_table = iommu_from_common(range->common);
	struct pt_state pts = pt_init(range, level, table);
	struct pt_iommu_map_args *map = arg;
	unsigned int leaf_pgsize_lg2 = map->leaf_pgsize_lg2;
	unsigned int leaves_avail;
	unsigned int start_index;
	pt_oaddr_t oa = map->oa;
	pt_vaddr_t num_leaves;
	unsigned int orig_end;
	unsigned int step_lg2;
	pt_vaddr_t last_va;
	unsigned int step;
	bool need_contig;
	int ret = 0;

	PT_WARN_ON(map->leaf_level != level);
	PT_WARN_ON(!pt_can_have_leaf(&pts));

	step_lg2 = leaf_pgsize_lg2 - pt_table_item_lg2sz(&pts);
	step = log2_to_int_t(unsigned int, step_lg2);
	need_contig = step_lg2 != 0;

	_pt_iter_first(&pts);
	start_index = pts.index;
	orig_end = pts.end_index;
	leaves_avail =
		log2_div_t(unsigned int, pts.end_index - pts.index, step_lg2);
	if (map->num_leaves <= leaves_avail) {
		/* Need to stop in the middle of the table to change sizes */
		pts.end_index = pts.index + log2_mul(map->num_leaves, step_lg2);
		num_leaves = 0;
	} else {
		num_leaves = map->num_leaves - leaves_avail;
	}

	PT_WARN_ON(
		log2_mod_t(unsigned int, pts.end_index - pts.index, step_lg2));
	do {
		pts.type = pt_load_entry_raw(&pts);
		if (pts.type != PT_ENTRY_EMPTY || need_contig) {
			if (pts.index != start_index)
				pt_index_to_va(&pts);
			ret = clear_contig(&pts, map->iotlb_gather, step,
					   leaf_pgsize_lg2);
			if (ret)
				break;
		}

		if (IS_ENABLED(CONFIG_DEBUG_GENERIC_PT)) {
			pt_index_to_va(&pts);
			PT_WARN_ON(compute_best_pgsize(&pts, oa) !=
				   leaf_pgsize_lg2);
		}
		pt_install_leaf_entry(&pts, oa, leaf_pgsize_lg2, &map->attrs);

		oa += log2_to_int(leaf_pgsize_lg2);
		pts.index += step;
	} while (pts.index < pts.end_index);

	flush_writes_range(&pts, start_index, pts.index);

	map->oa = oa;
	map->num_leaves = num_leaves;
	if (ret || num_leaves)
		return ret;

	/* range->va is not valid if we reached the end of the table */
	pts.index -= step;
	pt_index_to_va(&pts);
	pts.index += step;
	last_va = range->va + log2_to_int(leaf_pgsize_lg2);

	if (last_va - 1 == range->last_va) {
		PT_WARN_ON(pts.index != orig_end);
		return 0;
	}

	/*
	 * Reached a point where the page size changed, compute the new
	 * parameters.
	 */
	map->leaf_pgsize_lg2 = pt_compute_best_pgsize(
		iommu_table->domain.pgsize_bitmap, last_va, range->last_va, oa);
	map->leaf_level =
		pt_pgsz_lg2_to_level(range->common, map->leaf_pgsize_lg2);
	map->num_leaves = pt_pgsz_count(iommu_table->domain.pgsize_bitmap,
					last_va, range->last_va, oa,
					map->leaf_pgsize_lg2);

	/* Didn't finish this table level, caller will repeat it */
	if (pts.index != orig_end) {
		if (pts.index != start_index)
			pt_index_to_va(&pts);
		return -EAGAIN;
	}
	return 0;
}

static int __map_range(struct pt_range *range, void *arg, unsigned int level,
		       struct pt_table_p *table)
{
	struct pt_state pts = pt_init(range, level, table);
	struct pt_iommu_map_args *map = arg;
	int ret;

	PT_WARN_ON(map->leaf_level == level);
	PT_WARN_ON(!pt_can_have_table(&pts));

	_pt_iter_first(&pts);

	/* Descend to a child table */
	do {
		pts.type = pt_load_entry_raw(&pts);

		if (pts.type != PT_ENTRY_TABLE) {
			if (pts.type != PT_ENTRY_EMPTY)
				return -EADDRINUSE;
			ret = pt_iommu_new_table(&pts, &map->attrs);
			/* EAGAIN on a race will loop again */
			if (ret)
				return ret;
		} else {
			pts.table_lower = pt_table_ptr(&pts);
			/*
			 * Racing with a shared pt_iommu_new_table()? The other
			 * thread is still flushing the cache, so we have to
			 * also flush it to ensure that when our thread's map
			 * completes all the table items leading to our mapping
			 * are visible.
			 *
			 * This requires the pt_set_bit_release() to be a
			 * release of the cache flush so that this can acquire
			 * visibility at the iommu.
			 */
			if (pts_feature(&pts, PT_FEAT_DMA_INCOHERENT) &&
			    !pt_test_sw_bit_acquire(&pts,
						    SW_BIT_CACHE_FLUSH_DONE))
				flush_writes_item(&pts);
		}

		/*
		 * The already present table can possibly be shared with another
		 * concurrent map.
		 */
		do {
			if (map->leaf_level == level - 1)
				ret = pt_descend(&pts, arg, __map_range_leaf);
			else
				ret = pt_descend(&pts, arg, __map_range);
		} while (ret == -EAGAIN);
		if (ret)
			return ret;

		pts.index++;
		pt_index_to_va(&pts);
		if (pts.index >= pts.end_index)
			break;

		/*
		 * This level is currently running __map_range_leaf() which is
		 * not correct if the target level has been updated to this
		 * level. Have the caller invoke __map_range_leaf.
		 */
		if (map->leaf_level == level)
			return -EAGAIN;
	} while (true);
	return 0;
}

/*
 * Fast path for the easy case of mapping a 4k page to an already allocated
 * table. This is a common workload. If it returns EAGAIN run the full algorithm
 * instead.
 */
static __always_inline int __do_map_single_page(struct pt_range *range,
						void *arg, unsigned int level,
						struct pt_table_p *table,
						pt_level_fn_t descend_fn)
{
	struct pt_state pts = pt_init(range, level, table);
	struct pt_iommu_map_args *map = arg;

	pts.type = pt_load_single_entry(&pts);
	if (pts.level == 0) {
		if (pts.type != PT_ENTRY_EMPTY)
			return -EADDRINUSE;
		pt_install_leaf_entry(&pts, map->oa, PAGE_SHIFT,
				      &map->attrs);
		/* No flush, not used when incoherent */
		map->oa += PAGE_SIZE;
		return 0;
	}
	if (pts.type == PT_ENTRY_TABLE)
		return pt_descend(&pts, arg, descend_fn);
	/* Something else, use the slow path */
	return -EAGAIN;
}
PT_MAKE_LEVELS(__map_single_page, __do_map_single_page);

/*
 * Add a table to the top, increasing the top level as much as necessary to
 * encompass range.
 */
static int increase_top(struct pt_iommu *iommu_table, struct pt_range *range,
			struct pt_iommu_map_args *map)
{
	struct iommu_pages_list free_list = IOMMU_PAGES_LIST_INIT(free_list);
	struct pt_common *common = common_from_iommu(iommu_table);
	uintptr_t top_of_table = READ_ONCE(common->top_of_table);
	uintptr_t new_top_of_table = top_of_table;
	struct pt_table_p *table_mem;
	unsigned int new_level;
	spinlock_t *domain_lock;
	unsigned long flags;
	int ret;

	while (true) {
		struct pt_range top_range =
			_pt_top_range(common, new_top_of_table);
		struct pt_state pts = pt_init_top(&top_range);

		top_range.va = range->va;
		top_range.last_va = range->last_va;

		if (!pt_check_range(&top_range) &&
		    map->leaf_level <= pts.level) {
			new_level = pts.level;
			break;
		}

		pts.level++;
		if (pts.level > PT_MAX_TOP_LEVEL ||
		    pt_table_item_lg2sz(&pts) >= common->max_vasz_lg2) {
			ret = -ERANGE;
			goto err_free;
		}

		table_mem =
			table_alloc_top(common, _pt_top_set(NULL, pts.level),
					map->attrs.gfp, ALLOC_DEFER_COHERENT_FLUSH);
		if (IS_ERR(table_mem)) {
			ret = PTR_ERR(table_mem);
			goto err_free;
		}
		iommu_pages_list_add(&free_list, table_mem);

		/* The new table links to the lower table always at index 0 */
		top_range.va = 0;
		top_range.top_level = pts.level;
		pts.table_lower = pts.table;
		pts.table = table_mem;
		pt_load_single_entry(&pts);
		PT_WARN_ON(pts.index != 0);
		pt_install_table(&pts, virt_to_phys(pts.table_lower),
				 &map->attrs);
		new_top_of_table = _pt_top_set(pts.table, pts.level);
	}

	/*
	 * Avoid double flushing, flush it once after all pt_install_table()
	 */
	if (pt_feature(common, PT_FEAT_DMA_INCOHERENT)) {
		ret = iommu_pages_start_incoherent_list(
			&free_list, iommu_table->iommu_device);
		if (ret)
			goto err_free;
	}

	/*
	 * top_of_table is write locked by the spinlock, but readers can use
	 * READ_ONCE() to get the value. Since we encode both the level and the
	 * pointer in one quanta the lockless reader will always see something
	 * valid. The HW must be updated to the new level under the spinlock
	 * before top_of_table is updated so that concurrent readers don't map
	 * into the new level until it is fully functional. If another thread
	 * already updated it while we were working then throw everything away
	 * and try again.
	 */
	domain_lock = iommu_table->driver_ops->get_top_lock(iommu_table);
	spin_lock_irqsave(domain_lock, flags);
	if (common->top_of_table != top_of_table ||
	    top_of_table == new_top_of_table) {
		spin_unlock_irqrestore(domain_lock, flags);
		ret = -EAGAIN;
		goto err_free;
	}

	/*
	 * We do not issue any flushes for change_top on the expectation that
	 * any walk cache will not become a problem by adding another layer to
	 * the tree. Misses will rewalk from the updated top pointer, hits
	 * continue to be correct. Negative caching is fine too since all the
	 * new IOVA added by the new top is non-present.
	 */
	iommu_table->driver_ops->change_top(
		iommu_table, virt_to_phys(table_mem), new_level);
	WRITE_ONCE(common->top_of_table, new_top_of_table);
	spin_unlock_irqrestore(domain_lock, flags);
	return 0;

err_free:
	if (pt_feature(common, PT_FEAT_DMA_INCOHERENT))
		iommu_pages_stop_incoherent_list(&free_list,
						 iommu_table->iommu_device);
	iommu_put_pages_list(&free_list);
	return ret;
}

static int check_map_range(struct pt_iommu *iommu_table, struct pt_range *range,
			   struct pt_iommu_map_args *map)
{
	struct pt_common *common = common_from_iommu(iommu_table);
	int ret;

	do {
		ret = pt_check_range(range);
		if (!pt_feature(common, PT_FEAT_DYNAMIC_TOP))
			return ret;

		if (!ret && map->leaf_level <= range->top_level)
			break;

		ret = increase_top(iommu_table, range, map);
		if (ret && ret != -EAGAIN)
			return ret;

		/* Reload the new top */
		*range = pt_make_range(common, range->va, range->last_va);
	} while (ret);
	PT_WARN_ON(pt_check_range(range));
	return 0;
}

static int do_map(struct pt_range *range, struct pt_common *common,
		  bool single_page, struct pt_iommu_map_args *map)
{
	int ret;

	/*
	 * The __map_single_page() fast path does not support DMA_INCOHERENT
	 * flushing to keep its .text small.
	 */
	if (single_page && !pt_feature(common, PT_FEAT_DMA_INCOHERENT)) {

		ret = pt_walk_range(range, __map_single_page, map);
		if (ret != -EAGAIN)
			return ret;
		/* EAGAIN falls through to the full path */
	}

	do {
		if (map->leaf_level == range->top_level)
			ret = pt_walk_range(range, __map_range_leaf, map);
		else
			ret = pt_walk_range(range, __map_range, map);
	} while (ret == -EAGAIN);
	return ret;
}

static int NS(map_range)(struct pt_iommu *iommu_table, dma_addr_t iova,
			 phys_addr_t paddr, dma_addr_t len, unsigned int prot,
			 gfp_t gfp, size_t *mapped)
{
	pt_vaddr_t pgsize_bitmap = iommu_table->domain.pgsize_bitmap;
	struct pt_common *common = common_from_iommu(iommu_table);
	struct iommu_iotlb_gather iotlb_gather;
	struct pt_iommu_map_args map = {
		.iotlb_gather = &iotlb_gather,
		.oa = paddr,
	};
	bool single_page = false;
	struct pt_range range;
	int ret;

	iommu_iotlb_gather_init(&iotlb_gather);

	if (WARN_ON(!(prot & (IOMMU_READ | IOMMU_WRITE))))
		return -EINVAL;

	/* Check the paddr doesn't exceed what the table can store */
	if ((sizeof(pt_oaddr_t) < sizeof(paddr) &&
	     (pt_vaddr_t)paddr > PT_VADDR_MAX) ||
	    (common->max_oasz_lg2 != PT_VADDR_MAX_LG2 &&
	     oalog2_div(paddr, common->max_oasz_lg2)))
		return -ERANGE;

	ret = pt_iommu_set_prot(common, &map.attrs, prot);
	if (ret)
		return ret;
	map.attrs.gfp = gfp;

	ret = make_range_no_check(common, &range, iova, len);
	if (ret)
		return ret;

	/* Calculate target page size and level for the leaves */
	if (pt_has_system_page_size(common) && len == PAGE_SIZE &&
		likely(pgsize_bitmap & PAGE_SIZE)) {
		if (log2_mod(iova | paddr, PAGE_SHIFT))
			return -ENXIO;
		map.leaf_pgsize_lg2 = PAGE_SHIFT;
		map.leaf_level = 0;
		map.num_leaves = 1;
		single_page = true;
	} else {
		map.leaf_pgsize_lg2 = pt_compute_best_pgsize(
			pgsize_bitmap, range.va, range.last_va, paddr);
		if (!map.leaf_pgsize_lg2)
			return -ENXIO;
		map.leaf_level =
			pt_pgsz_lg2_to_level(common, map.leaf_pgsize_lg2);
		map.num_leaves = pt_pgsz_count(pgsize_bitmap, range.va,
					       range.last_va, paddr,
					       map.leaf_pgsize_lg2);
	}

	ret = check_map_range(iommu_table, &range, &map);
	if (ret)
		return ret;

	PT_WARN_ON(map.leaf_level > range.top_level);

	ret = do_map(&range, common, single_page, &map);

	/*
	 * Table levels were freed and replaced with large items, flush any walk
	 * cache that may refer to the freed levels.
	 */
	if (!iommu_pages_list_empty(&iotlb_gather.freelist))
		iommu_iotlb_sync(&iommu_table->domain, &iotlb_gather);

	/* Bytes successfully mapped */
	PT_WARN_ON(!ret && map.oa - paddr != len);
	*mapped += map.oa - paddr;
	return ret;
}

struct pt_unmap_args {
	struct iommupt_pending_gather pending;
	pt_vaddr_t unmapped;
};

static __maybe_unused int __unmap_range(struct pt_range *range, void *arg,
					unsigned int level,
					struct pt_table_p *table)
{
	struct pt_state pts = pt_init(range, level, table);
	unsigned int flush_start_index = UINT_MAX;
	unsigned int flush_end_index = UINT_MAX;
	struct pt_unmap_args *unmap = arg;
	unsigned int num_oas = 0;
	unsigned int start_index;
	int ret = 0;

	_pt_iter_first(&pts);
	start_index = pts.index;
	pts.type = pt_load_entry_raw(&pts);
	/*
	 * A starting index is in the middle of a contiguous entry
	 *
	 * The IOMMU API does not require drivers to support unmapping parts of
	 * large pages. Long ago VFIO would try to split maps but the current
	 * version never does.
	 *
	 * Instead when unmap reaches a partial unmap of the start of a large
	 * IOPTE it should remove the entire IOPTE and return that size to the
	 * caller.
	 */
	if (pts.type == PT_ENTRY_OA) {
		if (log2_mod(range->va, pt_entry_oa_lg2sz(&pts)))
			return -EINVAL;
		/* Micro optimization */
		goto start_oa;
	}

	do {
		if (pts.type != PT_ENTRY_OA) {
			bool fully_covered;

			if (pts.type != PT_ENTRY_TABLE) {
				ret = -EINVAL;
				break;
			}

			if (pts.index != start_index)
				pt_index_to_va(&pts);
			pts.table_lower = pt_table_ptr(&pts);

			fully_covered = pt_entry_fully_covered(
				&pts, pt_table_item_lg2sz(&pts));

			ret = pt_descend(&pts, arg, __unmap_range);
			if (ret)
				break;

			/*
			 * If the unmapping range fully covers the table then we
			 * can free it as well. The clear is delayed until we
			 * succeed in clearing the lower table levels.
			 */
			if (fully_covered) {
				gather_add_table(&unmap->pending, &pts,
						 pts.table_lower);
				pt_clear_entries(&pts, ilog2(1));
				if (pts.index < flush_start_index)
					flush_start_index = pts.index;
				flush_end_index = pts.index + 1;
			}
			pts.index++;
		} else {
			unsigned int num_contig_lg2;
start_oa:
			/*
			 * If the caller requested an last that falls within a
			 * single entry then the entire entry is unmapped and
			 * the length returned will be larger than requested.
			 */
			num_contig_lg2 = pt_entry_num_contig_lg2(&pts);
			pt_clear_entries(&pts, num_contig_lg2);
			gather_add_leaf(&unmap->pending, &pts);
			num_oas += log2_to_int(num_contig_lg2);
			if (pts.index < flush_start_index)
				flush_start_index = pts.index;
			pts.index += log2_to_int(num_contig_lg2);
			flush_end_index = pts.index;
		}
		if (pts.index >= pts.end_index)
			break;
		pts.type = pt_load_entry_raw(&pts);
	} while (true);

	unmap->unmapped += log2_mul(num_oas, pt_table_item_lg2sz(&pts));
	if (flush_start_index != flush_end_index)
		flush_writes_range(&pts, flush_start_index, flush_end_index);

	return ret;
}

static size_t NS(unmap_range)(struct pt_iommu *iommu_table, dma_addr_t iova,
			      dma_addr_t len,
			      struct iommu_iotlb_gather *iotlb_gather)
{
	struct pt_unmap_args unmap = {
		.pending.iotlb_gather = iotlb_gather,
		.pending.free_list = IOMMU_PAGES_LIST_INIT(
			unmap.pending.free_list),
	};
	struct pt_range range;
	int ret;

	ret = make_range(common_from_iommu(iommu_table), &range, iova, len);
	if (ret)
		return 0;

	pt_walk_range(&range, __unmap_range, &unmap);

	gather_range_pending(&unmap.pending, iommu_table, iova, unmap.unmapped);

	return unmap.unmapped;
}

static void NS(get_info)(struct pt_iommu *iommu_table,
			 struct pt_iommu_info *info)
{
	struct pt_common *common = common_from_iommu(iommu_table);
	struct pt_range range = pt_top_range(common);
	struct pt_state pts = pt_init_top(&range);
	pt_vaddr_t pgsize_bitmap = 0;

	if (pt_feature(common, PT_FEAT_DYNAMIC_TOP)) {
		for (pts.level = 0; pts.level <= PT_MAX_TOP_LEVEL;
		     pts.level++) {
			if (pt_table_item_lg2sz(&pts) >= common->max_vasz_lg2)
				break;
			pgsize_bitmap |= pt_possible_sizes(&pts);
		}
	} else {
		for (pts.level = 0; pts.level <= range.top_level; pts.level++)
			pgsize_bitmap |= pt_possible_sizes(&pts);
	}

	/*
	 * Hide page sizes larger than the maximum. -1 because a whole table
	 * pgsize is not allowed
	 */
	info->pgsize_bitmap = log2_mod(pgsize_bitmap, common->max_vasz_lg2 - 1);
	info->pgsize_bitmap = oalog2_mod(info->pgsize_bitmap, common->max_oasz_lg2);
}

static void NS(deinit)(struct pt_iommu *iommu_table)
{
	struct pt_common *common = common_from_iommu(iommu_table);
	struct pt_range range = pt_all_range(common);
	struct pt_iommu_collect_args collect = {
		.pending.free_list = IOMMU_PAGES_LIST_INIT(
			collect.pending.free_list),
	};

	iommu_pages_list_add(&collect.pending.free_list, range.top_table);
	pt_walk_range(&range, __collect_tables, &collect);

	/*
	 * The driver has to already have fenced the HW access to the page table
	 * and invalidated any caching referring to this memory.
	 */
	if (pt_feature(common, PT_FEAT_DMA_INCOHERENT))
		iommu_pages_stop_incoherent_list(&collect.pending.free_list,
						 iommu_table->iommu_device);
	iommu_put_pages_list(&collect.pending.free_list);
}

static const struct pt_iommu_ops NS(ops) = {
	.map_range = NS(map_range),
	.unmap_range = NS(unmap_range),
#if IS_ENABLED(CONFIG_IOMMUFD_DRIVER) && defined(pt_entry_is_write_dirty) && \
	IS_ENABLED(CONFIG_IOMMUFD_TEST) && defined(pt_entry_make_write_dirty)
	.set_dirty = NS(set_dirty),
#endif
	.get_info = NS(get_info),
	.deinit = NS(deinit),
};

static int pt_init_common(struct pt_common *common)
{
	struct pt_range top_range = pt_top_range(common);

	if (PT_WARN_ON(top_range.top_level > PT_MAX_TOP_LEVEL))
		return -EINVAL;

	if (top_range.max_vasz_lg2 == PT_VADDR_MAX_LG2)
		common->features |= BIT(PT_FEAT_FULL_VA);

	/* Requested features must match features compiled into this format */
	if ((common->features & ~(unsigned int)PT_SUPPORTED_FEATURES) ||
	    (!IS_ENABLED(CONFIG_DEBUG_GENERIC_PT) &&
	     (common->features & PT_FORCE_ENABLED_FEATURES) !=
		     PT_FORCE_ENABLED_FEATURES))
		return -EOPNOTSUPP;

	/*
	 * Check if the top level of the page table is too small to hold the
	 * specified maxvasz.
	 */
	if (!pt_feature(common, PT_FEAT_DYNAMIC_TOP) &&
	    top_range.top_level != PT_MAX_TOP_LEVEL) {
		struct pt_state pts = { .range = &top_range,
					.level = top_range.top_level };

		if (common->max_vasz_lg2 >
		    pt_num_items_lg2(&pts) + pt_table_item_lg2sz(&pts))
			return -EOPNOTSUPP;
	}

	if (common->max_oasz_lg2 == 0)
		common->max_oasz_lg2 = pt_max_oa_lg2(common);
	else
		common->max_oasz_lg2 = min(common->max_oasz_lg2,
					   pt_max_oa_lg2(common));
	return 0;
}

static int pt_iommu_init_domain(struct pt_iommu *iommu_table,
				struct iommu_domain *domain)
{
	struct pt_common *common = common_from_iommu(iommu_table);
	struct pt_iommu_info info;
	struct pt_range range;

	NS(get_info)(iommu_table, &info);

	domain->type = __IOMMU_DOMAIN_PAGING;
	domain->pgsize_bitmap = info.pgsize_bitmap;
	domain->is_iommupt = true;

	if (pt_feature(common, PT_FEAT_DYNAMIC_TOP))
		range = _pt_top_range(common,
				      _pt_top_set(NULL, PT_MAX_TOP_LEVEL));
	else
		range = pt_top_range(common);

	/* A 64-bit high address space table on a 32-bit system cannot work. */
	domain->geometry.aperture_start = (unsigned long)range.va;
	if ((pt_vaddr_t)domain->geometry.aperture_start != range.va)
		return -EOVERFLOW;

	/*
	 * The aperture is limited to what the API can do after considering all
	 * the different types dma_addr_t/unsigned long/pt_vaddr_t that are used
	 * to store a VA. Set the aperture to something that is valid for all
	 * cases. Saturate instead of truncate the end if the types are smaller
	 * than the top range. aperture_end should be called aperture_last.
	 */
	domain->geometry.aperture_end = (unsigned long)range.last_va;
	if ((pt_vaddr_t)domain->geometry.aperture_end != range.last_va) {
		domain->geometry.aperture_end = ULONG_MAX;
		domain->pgsize_bitmap &= ULONG_MAX;
	}
	domain->geometry.force_aperture = true;

	return 0;
}

static void pt_iommu_zero(struct pt_iommu_table *fmt_table)
{
	struct pt_iommu *iommu_table = &fmt_table->iommu;
	struct pt_iommu cfg = *iommu_table;

	static_assert(offsetof(struct pt_iommu_table, iommu.domain) == 0);
	memset_after(fmt_table, 0, iommu.domain);

	/* The caller can initialize some of these values */
	iommu_table->iommu_device = cfg.iommu_device;
	iommu_table->driver_ops = cfg.driver_ops;
	iommu_table->nid = cfg.nid;
}

#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)
{
	struct pt_iommu *iommu_table = &fmt_table->iommu;
	struct pt_common *common = common_from_iommu(iommu_table);
	struct pt_table_p *table_mem;
	int ret;

	if (cfg->common.hw_max_vasz_lg2 > PT_MAX_VA_ADDRESS_LG2 ||
	    !cfg->common.hw_max_vasz_lg2 || !cfg->common.hw_max_oasz_lg2)
		return -EINVAL;

	pt_iommu_zero(fmt_table);
	common->features = cfg->common.features;
	common->max_vasz_lg2 = cfg->common.hw_max_vasz_lg2;
	common->max_oasz_lg2 = cfg->common.hw_max_oasz_lg2;
	ret = pt_iommu_fmt_init(fmt_table, cfg);
	if (ret)
		return ret;

	if (cfg->common.hw_max_oasz_lg2 > pt_max_oa_lg2(common))
		return -EINVAL;

	ret = pt_init_common(common);
	if (ret)
		return ret;

	if (pt_feature(common, PT_FEAT_DYNAMIC_TOP) &&
	    WARN_ON(!iommu_table->driver_ops ||
		    !iommu_table->driver_ops->change_top ||
		    !iommu_table->driver_ops->get_top_lock))
		return -EINVAL;

	if (pt_feature(common, PT_FEAT_SIGN_EXTEND) &&
	    (pt_feature(common, PT_FEAT_FULL_VA) ||
	     pt_feature(common, PT_FEAT_DYNAMIC_TOP)))
		return -EINVAL;

	if (pt_feature(common, PT_FEAT_DMA_INCOHERENT) &&
	    WARN_ON(!iommu_table->iommu_device))
		return -EINVAL;

	ret = pt_iommu_init_domain(iommu_table, &iommu_table->domain);
	if (ret)
		return ret;

	table_mem = table_alloc_top(common, common->top_of_table, gfp,
				    ALLOC_NORMAL);
	if (IS_ERR(table_mem))
		return PTR_ERR(table_mem);
	pt_top_set(common, table_mem, pt_top_get_level(common));

	/* Must be last, see pt_iommu_deinit() */
	iommu_table->ops = &NS(ops);
	return 0;
}
EXPORT_SYMBOL_NS_GPL(pt_iommu_init, "GENERIC_PT_IOMMU");

#ifdef pt_iommu_fmt_hw_info
#define pt_iommu_table_hw_info CONCATENATE(pt_iommu_table, _hw_info)
#define pt_iommu_hw_info CONCATENATE(CONCATENATE(pt_iommu_, PTPFX), hw_info)
void pt_iommu_hw_info(struct pt_iommu_table *fmt_table,
		      struct pt_iommu_table_hw_info *info)
{
	struct pt_iommu *iommu_table = &fmt_table->iommu;
	struct pt_common *common = common_from_iommu(iommu_table);
	struct pt_range top_range = pt_top_range(common);

	pt_iommu_fmt_hw_info(fmt_table, &top_range, info);
}
EXPORT_SYMBOL_NS_GPL(pt_iommu_hw_info, "GENERIC_PT_IOMMU");
#endif

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("IOMMU Page table implementation for " __stringify(PTPFX_RAW));
MODULE_IMPORT_NS("GENERIC_PT");
/* For iommu_dirty_bitmap_record() */
MODULE_IMPORT_NS("IOMMUFD");

#endif  /* __GENERIC_PT_IOMMU_PT_H */