/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2026 Ruslan Bukin <br@bsdpad.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/systm.h>
#include <sys/ktr.h>
#include <sys/lock.h>
#include <sys/mutex.h>

#include <vm/vm.h>
#include <vm/vm_param.h>
#include <vm/vm_page.h>

#include <riscv/iommu/iommu_pmap.h>

/*
 * Boundary values for the page table page index space:
 *
 * L3 pages: [0, NUL2E)
 * L2 pages: [NUL2E, NUL2E + NUL1E)
 * L1 pages: [NUL2E + NUL1E, NUL2E + NUL1E + NUL0E)
 *
 * Note that these ranges are used in both SV39 and SV48 mode.  In SV39 mode the
 * ranges are not fully populated since there are at most Ln_ENTRIES^2 L3 pages
 * in a set of page tables.
 */
#define	NUL0E		Ln_ENTRIES
#define	NUL1E		(Ln_ENTRIES * NUL0E)
#define	NUL2E		(Ln_ENTRIES * NUL1E)

#define	pmap_l1_pindex(v)	(NUL2E + ((v) >> L1_SHIFT))
#define	pmap_l2_pindex(v)	((v) >> L2_SHIFT)

#define	pmap_clear(pte)			pmap_store(pte, 0)
#define	pmap_clear_bits(pte, bits)	atomic_clear_64(pte, bits)
#define	pmap_load_store(pte, entry)	atomic_swap_64(pte, entry)
#define	pmap_load_clear(pte)		pmap_load_store(pte, 0)
#define	pmap_load(pte)			atomic_load_64(pte)
#define	pmap_store(pte, entry)		atomic_store_64(pte, entry)
#define	pmap_store_bits(pte, bits)	atomic_set_64(pte, bits)

#define	pmap_l0_index(va)	(((va) >> L0_SHIFT) & Ln_ADDR_MASK)
#define	pmap_l1_index(va)	(((va) >> L1_SHIFT) & Ln_ADDR_MASK)
#define	pmap_l2_index(va)	(((va) >> L2_SHIFT) & Ln_ADDR_MASK)
#define	pmap_l3_index(va)	(((va) >> L3_SHIFT) & Ln_ADDR_MASK)

#define	PTE_TO_PHYS(pte) \
    ((((pte) & ~PTE_HI_MASK) >> PTE_PPN0_S) * PAGE_SIZE)
#define	L2PTE_TO_PHYS(l2) \
    ((((l2) & ~PTE_HI_MASK) >> PTE_PPN1_S) << L2_SHIFT)
#define	L1PTE_TO_PHYS(l1) \
    ((((l1) & ~PTE_HI_MASK) >> PTE_PPN2_S) << L1_SHIFT)
#define PTE_TO_VM_PAGE(pte) PHYS_TO_VM_PAGE(PTE_TO_PHYS(pte))

/********************/
/* Inline functions */
/********************/

static __inline pd_entry_t *
pmap_l0(struct riscv_iommu_pmap *pmap, vm_offset_t va)
{
	KASSERT(pmap->pm_mode != PMAP_MODE_SV39,
	    ("%s: in SV39 mode", __func__));
	KASSERT(VIRT_IS_VALID(va),
	    ("%s: malformed virtual address %#lx", __func__, va));
	return (&pmap->pm_top[pmap_l0_index(va)]);
}

static __inline pd_entry_t *
pmap_l0_to_l1(struct riscv_iommu_pmap *pmap, pd_entry_t *l0, vm_offset_t va)
{
	vm_paddr_t phys;
	pd_entry_t *l1;

	KASSERT(pmap->pm_mode != PMAP_MODE_SV39,
	    ("%s: in SV39 mode", __func__));
	phys = PTE_TO_PHYS(pmap_load(l0));
	l1 = (pd_entry_t *)PHYS_TO_DMAP(phys);

	return (&l1[pmap_l1_index(va)]);
}

static __inline pd_entry_t *
pmap_l1(struct riscv_iommu_pmap *pmap, vm_offset_t va)
{
	pd_entry_t *l0;

	KASSERT(VIRT_IS_VALID(va),
	    ("%s: malformed virtual address %#lx", __func__, va));
	if (pmap->pm_mode == PMAP_MODE_SV39) {
		return (&pmap->pm_top[pmap_l1_index(va)]);
	} else {
		l0 = pmap_l0(pmap, va);
		if ((pmap_load(l0) & PTE_V) == 0)
			return (NULL);
		if ((pmap_load(l0) & PTE_RX) != 0)
			return (NULL);
		return (pmap_l0_to_l1(pmap, l0, va));
	}
}

static __inline pd_entry_t *
pmap_l1_to_l2(pd_entry_t *l1, vm_offset_t va)
{
	vm_paddr_t phys;
	pd_entry_t *l2;

	phys = PTE_TO_PHYS(pmap_load(l1));
	l2 = (pd_entry_t *)PHYS_TO_DMAP(phys);

	return (&l2[pmap_l2_index(va)]);
}

static __inline pd_entry_t *
pmap_l2(struct riscv_iommu_pmap *pmap, vm_offset_t va)
{
	pd_entry_t *l1;

	l1 = pmap_l1(pmap, va);
	if (l1 == NULL)
		return (NULL);
	if ((pmap_load(l1) & PTE_V) == 0)
		return (NULL);
	if ((pmap_load(l1) & PTE_RX) != 0)
		return (NULL);

	return (pmap_l1_to_l2(l1, va));
}

static __inline pt_entry_t *
pmap_l2_to_l3(pd_entry_t *l2, vm_offset_t va)
{
	vm_paddr_t phys;
	pt_entry_t *l3;

	phys = PTE_TO_PHYS(pmap_load(l2));
	l3 = (pd_entry_t *)PHYS_TO_DMAP(phys);

	return (&l3[pmap_l3_index(va)]);
}

static __inline pt_entry_t *
pmap_l3(struct riscv_iommu_pmap *pmap, vm_offset_t va)
{
	pd_entry_t *l2;

	l2 = pmap_l2(pmap, va);
	if (l2 == NULL)
		return (NULL);
	if ((pmap_load(l2) & PTE_V) == 0)
		return (NULL);
	if ((pmap_load(l2) & PTE_RX) != 0)
		return (NULL);

	return (pmap_l2_to_l3(l2, va));
}

#ifdef INVARIANTS
static __inline void
pmap_resident_count_inc(struct riscv_iommu_pmap *pmap, int count)
{

	PMAP_LOCK_ASSERT(pmap, MA_OWNED);
	pmap->sp_resident_count += count;
}

static __inline void
pmap_resident_count_dec(struct riscv_iommu_pmap *pmap, int count)
{

	PMAP_LOCK_ASSERT(pmap, MA_OWNED);
	KASSERT(pmap->sp_resident_count >= count,
	    ("pmap %p resident count underflow %ld %d", pmap,
	    pmap->sp_resident_count, count));
	pmap->sp_resident_count -= count;
}
#else
static __inline void
pmap_resident_count_inc(struct riscv_iommu_pmap *pmap, int count)
{
}

static __inline void
pmap_resident_count_dec(struct riscv_iommu_pmap *pmap, int count)
{
}
#endif

/***************************************************
 * Page table page management routines.....
 ***************************************************/

int
iommu_pmap_pinit(struct riscv_iommu_pmap *pmap, enum pmap_mode pm_mode)
{
	vm_paddr_t topphys;
	vm_page_t m;

	m = vm_page_alloc_noobj(VM_ALLOC_WIRED | VM_ALLOC_ZERO |
	    VM_ALLOC_WAITOK);
	topphys = VM_PAGE_TO_PHYS(m);
	pmap->pm_top = (pd_entry_t *)PHYS_TO_DMAP(topphys);
	pmap->pm_mode = pm_mode;

	switch (pm_mode) {
	case PMAP_MODE_SV39:
		pmap->pm_satp = SATP_MODE_SV39;
		break;
	case PMAP_MODE_SV48:
		pmap->pm_satp = SATP_MODE_SV48;
		break;
	default:
		panic("Unknown virtual memory system");
	};

	pmap->pm_satp |= (topphys >> PAGE_SHIFT);

#ifdef INVARIANTS
	pmap->sp_resident_count = 0;
#endif

	mtx_init(&pmap->pm_mtx, "iommu pmap", NULL, MTX_DEF);

	return (1);
}

/*
 * Release any resources held by the given physical map.
 * Called when a pmap initialized by pmap_pinit is being released.
 * Should only be called if the map contains no valid mappings.
 */
void
iommu_pmap_release(struct riscv_iommu_pmap *pmap)
{
	vm_page_t m;

	KASSERT(pmap->sp_resident_count == 0,
	    ("pmap_release: pmap resident count %ld != 0",
	    pmap->sp_resident_count));

	m = PHYS_TO_VM_PAGE(DMAP_TO_PHYS((vm_offset_t)pmap->pm_top));
	vm_page_unwire_noq(m);
	vm_page_free_zero(m);
	mtx_destroy(&pmap->pm_mtx);
}

/*
 * This routine is called if the desired page table page does not exist.
 *
 * If page table page allocation fails, this routine may sleep before
 * returning NULL.  It sleeps only if a lock pointer was given.
 *
 * Note: If a page allocation fails at page table level two or three,
 * one or two pages may be held during the wait, only to be released
 * afterwards.  This conservative approach is easily argued to avoid
 * race conditions.
 */
static vm_page_t
_pmap_alloc_l3(struct riscv_iommu_pmap *pmap, vm_pindex_t ptepindex)
{
	vm_page_t m, pdpg;
	pt_entry_t entry;
	vm_paddr_t phys;
	pn_t pn;

	PMAP_LOCK_ASSERT(pmap, MA_OWNED);

	/*
	 * Allocate a page table page.
	 */
	m = vm_page_alloc_noobj(VM_ALLOC_WIRED | VM_ALLOC_ZERO);
	if (m == NULL) {
		/*
		 * Indicate the need to retry.  While waiting, the page table
		 * page may have been allocated.
		 */
		return (NULL);
	}
	m->pindex = ptepindex;

	/*
	 * Map the pagetable page into the process address space, if
	 * it isn't already there.
	 */
	pn = VM_PAGE_TO_PHYS(m) >> PAGE_SHIFT;
	if (ptepindex >= NUL2E + NUL1E) {
		pd_entry_t *l0;
		vm_pindex_t l0index;

		KASSERT(pmap->pm_mode != PMAP_MODE_SV39,
		    ("%s: pindex %#lx in SV39 mode", __func__, ptepindex));
		KASSERT(ptepindex < NUL2E + NUL1E + NUL0E,
		    ("%s: pindex %#lx out of range", __func__, ptepindex));

		l0index = ptepindex - (NUL2E + NUL1E);
		l0 = &pmap->pm_top[l0index];
		KASSERT((pmap_load(l0) & PTE_V) == 0,
		    ("%s: L0 entry %#lx is valid", __func__, pmap_load(l0)));

		entry = PTE_V | (pn << PTE_PPN0_S);
		pmap_store(l0, entry);
	} else if (ptepindex >= NUL2E) {
		pd_entry_t *l0, *l1;
		vm_pindex_t l0index, l1index;

		l1index = ptepindex - NUL2E;
		if (pmap->pm_mode == PMAP_MODE_SV39) {
			l1 = &pmap->pm_top[l1index];
		} else {
			l0index = l1index >> Ln_ENTRIES_SHIFT;
			l0 = &pmap->pm_top[l0index];
			if (pmap_load(l0) == 0) {
				/* Recurse to allocate the L1 page. */
				if (_pmap_alloc_l3(pmap,
				    NUL2E + NUL1E + l0index) == NULL)
					goto fail;
				phys = PTE_TO_PHYS(pmap_load(l0));
			} else {
				phys = PTE_TO_PHYS(pmap_load(l0));
				pdpg = PHYS_TO_VM_PAGE(phys);
				pdpg->ref_count++;
			}
			l1 = (pd_entry_t *)PHYS_TO_DMAP(phys);
			l1 = &l1[ptepindex & Ln_ADDR_MASK];
		}
		KASSERT((pmap_load(l1) & PTE_V) == 0,
		    ("%s: L1 entry %#lx is valid", __func__, pmap_load(l1)));

		entry = PTE_V | (pn << PTE_PPN0_S);
		pmap_store(l1, entry);
	} else {
		vm_pindex_t l0index, l1index;
		pd_entry_t *l0, *l1, *l2;

		l1index = ptepindex >> (L1_SHIFT - L2_SHIFT);
		if (pmap->pm_mode == PMAP_MODE_SV39) {
			l1 = &pmap->pm_top[l1index];
			if (pmap_load(l1) == 0) {
				/* recurse for allocating page dir */
				if (_pmap_alloc_l3(pmap, NUL2E + l1index)
				    == NULL)
					goto fail;
			} else {
				pdpg = PTE_TO_VM_PAGE(pmap_load(l1));
				pdpg->ref_count++;
			}
		} else {
			l0index = l1index >> Ln_ENTRIES_SHIFT;
			l0 = &pmap->pm_top[l0index];
			if (pmap_load(l0) == 0) {
				/* Recurse to allocate the L1 entry. */
				if (_pmap_alloc_l3(pmap, NUL2E + l1index)
				    == NULL)
					goto fail;
				phys = PTE_TO_PHYS(pmap_load(l0));
				l1 = (pd_entry_t *)PHYS_TO_DMAP(phys);
				l1 = &l1[l1index & Ln_ADDR_MASK];
			} else {
				phys = PTE_TO_PHYS(pmap_load(l0));
				l1 = (pd_entry_t *)PHYS_TO_DMAP(phys);
				l1 = &l1[l1index & Ln_ADDR_MASK];
				if (pmap_load(l1) == 0) {
					/* Recurse to allocate the L2 page. */
					if (_pmap_alloc_l3(pmap,
					    NUL2E + l1index) == NULL)
						goto fail;
				} else {
					pdpg = PTE_TO_VM_PAGE(pmap_load(l1));
					pdpg->ref_count++;
				}
			}
		}

		phys = PTE_TO_PHYS(pmap_load(l1));
		l2 = (pd_entry_t *)PHYS_TO_DMAP(phys);
		l2 = &l2[ptepindex & Ln_ADDR_MASK];
		KASSERT((pmap_load(l2) & PTE_V) == 0,
		    ("%s: L2 entry %#lx is valid", __func__, pmap_load(l2)));

		entry = PTE_V | (pn << PTE_PPN0_S);
		pmap_store(l2, entry);
	}

	pmap_resident_count_inc(pmap, 1);

	return (m);

fail:
	vm_page_unwire_noq(m);
	vm_page_free_zero(m);
	return (NULL);
}

/*
 * Remove a single IOMMU entry.
 */
int
iommu_pmap_remove(struct riscv_iommu_pmap *pmap, vm_offset_t va)
{
	pt_entry_t *l3;
	int rc;

	PMAP_LOCK(pmap);

	l3 = pmap_l3(pmap, va);
	if (l3 != NULL) {
		pmap_resident_count_dec(pmap, 1);
		pmap_clear(l3);
		rc = KERN_SUCCESS;
	} else
		rc = KERN_FAILURE;

	PMAP_UNLOCK(pmap);

	return (rc);
}

/* Add a single IOMMU entry. This function does not sleep. */
int
iommu_pmap_enter(struct riscv_iommu_pmap *pmap, vm_offset_t va, vm_paddr_t pa,
    vm_prot_t prot, u_int flags)
{
	pd_entry_t *l2, l2e;
	pt_entry_t new_l3;
	pt_entry_t *l3;
	vm_page_t mpte;
	pn_t pn;
	int rv;

	pn = (pa / PAGE_SIZE);

	new_l3 = PTE_V | PTE_R | PTE_A;
	if (prot & VM_PROT_EXECUTE)
		new_l3 |= PTE_X;
	if (flags & VM_PROT_WRITE)
		new_l3 |= PTE_D;
	if (prot & VM_PROT_WRITE)
		new_l3 |= PTE_W;
	if (va < VM_MAX_USER_ADDRESS)
		new_l3 |= PTE_U;

	new_l3 |= (pn << PTE_PPN0_S);
	new_l3 |= PTE_MA_IO;

	/*
	 * Set modified bit gratuitously for writeable mappings if
	 * the page is unmanaged. We do not want to take a fault
	 * to do the dirty bit accounting for these mappings.
	 */
	if (prot & VM_PROT_WRITE)
		new_l3 |= PTE_D;

	CTR2(KTR_PMAP, "pmap_enter: %.16lx -> %.16lx", va, pa);

	mpte = NULL;
	PMAP_LOCK(pmap);

	l2 = pmap_l2(pmap, va);
	if (l2 != NULL && ((l2e = pmap_load(l2)) & PTE_V) != 0 &&
	    ((l2e & PTE_RWX) == 0)) {
		l3 = pmap_l2_to_l3(l2, va);
	} else if (va < VM_MAXUSER_ADDRESS) {
		mpte = _pmap_alloc_l3(pmap, pmap_l2_pindex(va));
		if (mpte == NULL) {
			CTR0(KTR_PMAP, "pmap_enter: mpte == NULL");
			rv = KERN_RESOURCE_SHORTAGE;
			goto out;
		}
		l3 = pmap_l3(pmap, va);
	} else
		panic("pmap_enter: missing L3 table for kernel va %#lx", va);

	KASSERT((pmap_load(l3) & PTE_V) == 0, ("l3 is valid"));

	pmap_store(l3, new_l3);
	pmap_resident_count_inc(pmap, 1);

	rv = KERN_SUCCESS;
out:
	PMAP_UNLOCK(pmap);

	return (rv);
}

static void
iommu_pmap_remove_pages_sv48(struct riscv_iommu_pmap *pmap)
{
	pd_entry_t l0e, *l1, l1e, *l2, l2e, *l3, l3e;
	vm_paddr_t pa0, pa1, pa;
	vm_page_t m0, m1, m;
	int i, j, k, l;

	PMAP_LOCK(pmap);

	for (i = 0; i < Ln_ENTRIES; i++) {
		l0e = pmap->pm_top[i];
		if ((l0e & PTE_V) == 0)
			continue;
		pa0 = PTE_TO_PHYS(l0e);
		m0 = PHYS_TO_VM_PAGE(pa0);
		l1 = (pd_entry_t *)PHYS_TO_DMAP(pa0);

		for (j = 0; j < Ln_ENTRIES; j++) {
			l1e = l1[j];
			if ((l1e & PTE_V) == 0)
				continue;
			pa1 = PTE_TO_PHYS(l1e);
			m1 = PHYS_TO_VM_PAGE(pa1);
			l2 = (pd_entry_t *)PHYS_TO_DMAP(pa1);

			for (k = 0; k < Ln_ENTRIES; k++) {
				l2e = l2[k];
				if ((l2e & PTE_V) == 0)
					continue;
				pa = PTE_TO_PHYS(l2e);
				m = PHYS_TO_VM_PAGE(pa);
				l3 = (pt_entry_t *)PHYS_TO_DMAP(pa);

				for (l = 0; l < Ln_ENTRIES; l++) {
					l3e = l3[l];
					if ((l3e & PTE_V) == 0)
						continue;
					panic("%s: l3e found (idx %d %d %d %d)",
					    __func__, i, j, k, l);
				}

				vm_page_unwire_noq(m1);
				vm_page_unwire_noq(m);
				pmap_resident_count_dec(pmap, 1);
				vm_page_free(m);
				pmap_clear(&l2[k]);
			}

			vm_page_unwire_noq(m0);
			pmap_resident_count_dec(pmap, 1);
			vm_page_free(m1);
			pmap_clear(&l1[j]);
		}

		pmap_resident_count_dec(pmap, 1);
		vm_page_free(m0);
		pmap_clear(&pmap->pm_top[i]);
	}

	KASSERT(pmap->sp_resident_count == 0,
	    ("Invalid resident count %jd", pmap->sp_resident_count));

	PMAP_UNLOCK(pmap);
}

static void
iommu_pmap_remove_pages_sv39(struct riscv_iommu_pmap *pmap)
{
	pd_entry_t l1e, *l2, l2e, *l3, l3e;
	vm_paddr_t pa1, pa;
	vm_page_t m1, m;
	int j, k, l;

	PMAP_LOCK(pmap);

	for (j = 0; j < Ln_ENTRIES; j++) {
		l1e = pmap->pm_top[j];
		if ((l1e & PTE_V) == 0)
			continue;
		pa1 = PTE_TO_PHYS(l1e);
		m1 = PHYS_TO_VM_PAGE(pa1);
		l2 = (pd_entry_t *)PHYS_TO_DMAP(pa1);

		for (k = 0; k < Ln_ENTRIES; k++) {
			l2e = l2[k];
			if ((l2e & PTE_V) == 0)
				continue;
			pa = PTE_TO_PHYS(l2e);
			m = PHYS_TO_VM_PAGE(pa);
			l3 = (pt_entry_t *)PHYS_TO_DMAP(pa);

			for (l = 0; l < Ln_ENTRIES; l++) {
				l3e = l3[l];
				if ((l3e & PTE_V) == 0)
					continue;
				panic("%s: l3e found (idx %d %d %d)",
				    __func__, j, k, l);
			}

			vm_page_unwire_noq(m1);
			vm_page_unwire_noq(m);
			pmap_resident_count_dec(pmap, 1);
			vm_page_free(m);
			pmap_clear(&l2[k]);
		}

		pmap_resident_count_dec(pmap, 1);
		vm_page_free(m1);
		pmap_clear(&pmap->pm_top[j]);
	}

	KASSERT(pmap->sp_resident_count == 0,
	    ("Invalid resident count %jd", pmap->sp_resident_count));

	PMAP_UNLOCK(pmap);
}

void
iommu_pmap_remove_pages(struct riscv_iommu_pmap *pmap)
{

	switch (pmap->pm_mode) {
	case PMAP_MODE_SV39:
		iommu_pmap_remove_pages_sv39(pmap);
		break;
	case PMAP_MODE_SV48:
		iommu_pmap_remove_pages_sv48(pmap);
		break;
	default:
		panic("Unknown virtual memory system");
	}
}