/*-
 * 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/bitstring.h>
#include <sys/bus.h>
#include <sys/kernel.h>
#include <sys/malloc.h>
#include <sys/mutex.h>
#include <sys/rman.h>
#include <sys/lock.h>
#include <sys/sysctl.h>
#include <sys/tree.h>
#include <sys/taskqueue.h>
#include <sys/refcount.h>
#include <vm/vm.h>
#include <vm/vm_page.h>

#include <sys/bus.h>
#include <sys/conf.h>
#include <sys/rman.h>

#include <dev/pci/pcireg.h>
#include <dev/pci/pcivar.h>

#include <machine/bus.h>
#include <machine/resource.h>

#include <dev/iommu/iommu.h>
#include <riscv/iommu/iommu_pmap.h>
#include <riscv/iommu/iommu.h>

#include "iommu_if.h"

#define	dprintf(fmt, ...)

MALLOC_DEFINE(M_IOMMU, "RISCV_IOMMU", "RISC-V IOMMU");

#define	RD4(sc, reg)		bus_read_4(sc->res[0], (reg))
#define	WR4(sc, reg, val)	bus_write_4(sc->res[0], (reg), (val))
#define	RD8(sc, reg)		bus_read_8(sc->res[0], (reg))
#define	WR8(sc, reg, val)	bus_write_8(sc->res[0], (reg), (val))

#define	CQ_ENTRY_DWORDS		2	/* 16-byte */
#define	CQ_ENTRY_COUNT		8192	/* Amount of 16-byte entries. */
#define	FQ_ENTRY_DWORDS		4	/* 32-byte */
#define	FQ_ENTRY_COUNT		8192	/* Amount of 32-byte entries. */
#define	PQ_ENTRY_DWORDS		2	/* 16-byte */
#define	PQ_ENTRY_COUNT		8192	/* Amount of 16-byte entries. */

#define	DDT_NON_LEAF_DWORDS	1
#define	DDT_DC_STD_DWORDS	4	/* Standard-format DC. */
#define	DDT_DC_EXT_DWORDS	8	/* Extended-format DC. */
#define	DDT_L1_DID_BITS		9	/* All formats. */

#define	QUEUE_ALIGN		(1024 * 1024)	/* TODO */
#define	QUEUE_HEAD(q)		((q)->csr + RISCV_IOMMU_CQH - RISCV_IOMMU_CQB)
#define	QUEUE_TAIL(q)		((q)->csr + RISCV_IOMMU_CQT - RISCV_IOMMU_CQB)
#define	QUEUE_IPSR(q)		(1 << (q)->idx)

#define	PHYS_TO_PPN(p)		((p) >> 12)

struct riscv_iommu_fq_event {
	uint16_t cause_id;
	char *descr;
};

static struct riscv_iommu_fq_event fq_events[] = {
	{ FQ_CAUSE_INST_FAULT,		"Instruction access fault" },
	{ FQ_CAUSE_RD_ADDR_MISALIGNED,	"Read address misaligned" },
	{ FQ_CAUSE_RD_FAULT,		"Read access fault" },
	{ FQ_CAUSE_WR_ADDR_MISALIGNED,	"Write/AMO address misaligned" },
	{ FQ_CAUSE_WR_FAULT,		"Write/AMO access fault" },
	{ FQ_CAUSE_INST_FAULT_S,	"Instruction page fault" },
	{ FQ_CAUSE_RD_FAULT_S,		"Read page fault" },
	{ FQ_CAUSE_WR_FAULT_S,		"Write/AMO page fault" },
	{ FQ_CAUSE_INST_FAULT_VS,	"Instruction guest page fault" },
	{ FQ_CAUSE_RD_FAULT_VS,		"Read guest-page fault" },
	{ FQ_CAUSE_WR_FAULT_VS,		"Write/AMO guest-page fault" },
	{ FQ_CAUSE_DMA_DISABLED,	"All inbound transactions disallowed" },
	{ FQ_CAUSE_DDT_LOAD_FAULT,	"DDT entry load access fault" },
	{ FQ_CAUSE_DDT_INVALID,		"DDT entry not valid" },
	{ FQ_CAUSE_DDT_MISCONFIGURED,	"DDT entry misconfigured" },
	{ FQ_CAUSE_TR_TYPE_DISALLOWED,	"Transaction type disallowed" },
	{ FQ_CAUSE_MSI_LOAD_FAULT,	"MSI PTE load access fault" },
	{ FQ_CAUSE_MSI_INVALID, 	"MSI PTE not valid" },
	{ FQ_CAUSE_MSI_MISCONFIGURED,	"MSI PTE misconfigured" },
	{ FQ_CAUSE_MRIF_FAULT,		"MRIF access fault" },
	{ FQ_CAUSE_PDT_LOAD_FAULT,	"PDT entry load access fault" },
	{ FQ_CAUSE_PDT_INVALID,		"PDT entry not valid" },
	{ FQ_CAUSE_PDT_MISCONFIGURED,	"PDT entry misconfigured" },
	{ FQ_CAUSE_DDT_CORRUPTED,	"DDT data corruption" },
	{ FQ_CAUSE_PDT_CORRUPTED,	"PDT data corruption" },
	{ FQ_CAUSE_MSI_PT_CORRUPTED,	"MSI PT data corruption" },
	{ FQ_CAUSE_MRIF_CORRUPTED,	"MSI MRIF data corruption" },
	{ FQ_CAUSE_INTERNAL_DP_ERROR,	"Internal data path error" },
	{ FQ_CAUSE_MSI_WR_FAULT,	"IOMMU MSI write access fault" },
	{ FQ_CAUSE_PT_CORRUPTED,	"1st/2nd-stage PT data corruption" },
	{ 0, NULL },
};

static void
riscv_iommu_init_pscids(struct riscv_iommu_softc *sc)
{

	sc->pscid_set_size = (1 << sc->pscid_bits);
	sc->pscid_set = bit_alloc(sc->pscid_set_size, M_IOMMU, M_WAITOK);
	mtx_init(&sc->pscid_set_mutex, "pscid set", NULL, MTX_SPIN);
}

static int
riscv_iommu_pscid_alloc(struct riscv_iommu_softc *sc, int *new_pscid)
{

	mtx_lock_spin(&sc->pscid_set_mutex);
	bit_ffc(sc->pscid_set, sc->pscid_set_size, new_pscid);
	if (*new_pscid == -1) {
		mtx_unlock_spin(&sc->pscid_set_mutex);
		return (ENOMEM);
	}
	bit_set(sc->pscid_set, *new_pscid);
	mtx_unlock_spin(&sc->pscid_set_mutex);

	return (0);
}

static void
riscv_iommu_pscid_free(struct riscv_iommu_softc *sc, int pscid)
{

	mtx_lock_spin(&sc->pscid_set_mutex);
	bit_clear(sc->pscid_set, pscid);
	mtx_unlock_spin(&sc->pscid_set_mutex);
}

static uint32_t
riscv_iommu_q_inc_tail(struct riscv_iommu_queue *q)
{

	return ((q->lc.tail + 1) & q->mask);
}

static uint32_t
riscv_iommu_q_inc_head(struct riscv_iommu_queue *q)
{

	return ((q->lc.head + 1) & q->mask);
}

static int
riscv_iommu_q_has_space(struct riscv_iommu_queue *q)
{

	if (riscv_iommu_q_inc_tail(q) != q->lc.head)
		return (1);

	return (0);
}

static int
riscv_iommu_q_empty(struct riscv_iommu_queue *q)
{

	if (q->lc.tail == q->lc.head)
		return (1);

	return (0);
}

static int
riscv_iommu_dequeue(struct riscv_iommu_softc *sc, struct riscv_iommu_queue *q,
    void *data)
{
	void *entry_addr;

	q->lc.val = RD8(sc, q->head_off);
	entry_addr = (void *)((uint64_t)q->vaddr + q->lc.head * q->entry_size);
	memcpy(data, entry_addr, q->entry_size);
	q->lc.head = riscv_iommu_q_inc_head(q);
	WR4(sc, q->head_off, q->lc.head);

	return (0);
}

static int
riscv_iommu_enqueue(struct riscv_iommu_softc *sc, struct riscv_iommu_queue *q,
    void *data)
{
	void *entry_addr;

	RISCV_IOMMU_LOCK(sc);

	/* Ensure that a space is available. */
	do {
		q->lc.head = RD4(sc, q->head_off);
	} while (riscv_iommu_q_has_space(q) == 0);

	/* Write the command to the current tail entry. */
	entry_addr = (void *)((uint64_t)q->vaddr + q->lc.tail * q->entry_size);
	memcpy(entry_addr, data, q->entry_size);

	/* Increment tail index. */
	q->lc.tail = riscv_iommu_q_inc_tail(q);
	WR4(sc, q->tail_off, q->lc.tail);

	RISCV_IOMMU_UNLOCK(sc);

	return (0);
}

static void
riscv_iommu_sync(struct riscv_iommu_softc *sc, struct riscv_iommu_queue *q)
{
	struct riscv_iommu_command cmd;
	uint64_t reg;

	bzero(&cmd, sizeof(struct riscv_iommu_command));
	reg = COMMAND_OPCODE_IOFENCE;
	reg |= FUNC_IOFENCE_FUNC_C | FUNC_IOFENCE_PR | FUNC_IOFENCE_PW;
	cmd.dword0 = reg;

	riscv_iommu_enqueue(sc, &sc->cq, (void *)&cmd);

	/*
	 * FUNC_IOFENCE_WSI does not seem to be implemented in QEMU,
	 * so ensure all requests are processed in polling mode;
	 */
	do {
		q->lc.head = RD4(sc, q->head_off);
	} while (riscv_iommu_q_empty(q) == 0);
}

static int
riscv_iommu_inval_ddt(struct riscv_iommu_softc *sc)
{
	struct riscv_iommu_command cmd;
	uint64_t reg;

	bzero(&cmd, sizeof(struct riscv_iommu_command));
	reg = COMMAND_OPCODE_IODIR;
	reg |= FUNC_IODIR_INVAL_DDT;
	cmd.dword0 = reg;

	riscv_iommu_enqueue(sc, &sc->cq, (void *)&cmd);

	return (0);
}

static int
riscv_iommu_inval_ddt_did(struct riscv_iommu_softc *sc, int did)
{
	struct riscv_iommu_command cmd;
	uint64_t reg;

	bzero(&cmd, sizeof(struct riscv_iommu_command));
	reg = COMMAND_OPCODE_IODIR;
	reg |= FUNC_IODIR_INVAL_DDT;
	reg |= FUNC_IODIR_DV;
	reg |= (uint64_t)did << FUNC_IODIR_DID_S;
	cmd.dword0 = reg;

	riscv_iommu_enqueue(sc, &sc->cq, (void *)&cmd);

	return (0);
}

/* Invalidate entire address space. */
static int
riscv_iommu_inval_vma(struct riscv_iommu_softc *sc)
{
	struct riscv_iommu_command cmd;
	uint64_t reg;

	bzero(&cmd, sizeof(struct riscv_iommu_command));
	reg = COMMAND_OPCODE_IOTINVAL;
	reg |= FUNC_IOTINVAL_VMA;
	cmd.dword0 = reg;

	riscv_iommu_enqueue(sc, &sc->cq, (void *)&cmd);

	return (0);
}

static int
riscv_iommu_inval_vma_page(struct riscv_iommu_softc *sc, vm_offset_t addr,
    int pscid)
{
	struct riscv_iommu_command cmd;
	uint64_t reg;

	bzero(&cmd, sizeof(struct riscv_iommu_command));
	reg = COMMAND_OPCODE_IOTINVAL;
	reg |= FUNC_IOTINVAL_VMA;
	reg |= FUNC_IOTINVAL_AV;
	reg |= FUNC_IOTINVAL_PSCV;
	reg |= pscid << FUNC_IOTINVAL_PSCID_S;
	cmd.dword0 = reg;
	cmd.dword1 = PHYS_TO_PPN(addr) << FUNC_IOTINVAL_ADDR_S;

	riscv_iommu_enqueue(sc, &sc->cq, (void *)&cmd);

	return (0);
}

static int
riscv_iommu_inval_vma_pscid(struct riscv_iommu_softc *sc, int pscid)
{
	struct riscv_iommu_command cmd;
	uint64_t reg;

	bzero(&cmd, sizeof(struct riscv_iommu_command));
	reg = COMMAND_OPCODE_IOTINVAL;
	reg |= FUNC_IOTINVAL_VMA;
	reg |= FUNC_IOTINVAL_PSCV;
	reg |= pscid << FUNC_IOTINVAL_PSCID_S;
	cmd.dword0 = reg;

	riscv_iommu_enqueue(sc, &sc->cq, (void *)&cmd);

	return (0);
}

static int
riscv_iommu_set_mode(struct riscv_iommu_softc *sc)
{
	struct riscv_iommu_ddt *ddt;
	uint64_t reg;
	uint64_t base;

	reg = RD8(sc, RISCV_IOMMU_DDTP);
	if (reg & DDTP_BUSY)
		return (ENXIO);

	ddt = &sc->ddt;
	base = ddt->base | (sc->iommu_mode << DDTP_IOMMU_MODE_S);
	WR8(sc, RISCV_IOMMU_DDTP, base);

	reg = RD8(sc, RISCV_IOMMU_DDTP);
	if (reg != base) {
		device_printf(sc->dev, "could not set mode\n");
		return (ENXIO);
	}

	riscv_iommu_inval_ddt(sc);
	riscv_iommu_inval_vma(sc);

	return (0);
}

static int
riscv_iommu_enable_queue(struct riscv_iommu_softc *sc,
    struct riscv_iommu_queue *q)
{
	uint32_t reg;
	int timeout;

	if (q == &sc->cq)
		WR4(sc, QUEUE_TAIL(q), 0);
	else
		WR4(sc, QUEUE_HEAD(q), 0);

	reg = CQCSR_CQEN | CQCSR_CIE | CQCSR_CQMF;
	WR4(sc, q->csr, reg);

	timeout = 1000;
	do {
		reg = RD4(sc, RISCV_IOMMU_CQCSR);
		if ((reg & CQCSR_BUSY) == 0)
			break;
		DELAY(10);
	} while (timeout--);

	if (timeout <= 0) {
		device_printf(sc->dev, "could not enable command queue\n");
		return (-1);
	}

	if ((reg & CQCSR_CQON) == 0) {
		device_printf(sc->dev, "could not activate command queue\n");
		return (-1);
	}

	/* RW1C interrupt pending bit. */
	WR4(sc, RISCV_IOMMU_IPSR, QUEUE_IPSR(q));

	return (0);
}

static int
riscv_iommu_init_queue(struct riscv_iommu_softc *sc,
    struct riscv_iommu_queue *q, uint64_t base, uint32_t dwords)
{
	uint64_t reg;
	int sz;

	q->entry_size = dwords * 8;
	sz = (1 << q->size_log2) * q->entry_size;

	/* Set up the command circular buffer */
	q->vaddr = contigmalloc(sz, M_IOMMU, M_WAITOK | M_ZERO, 0,
	    (1ul << 48) - 1, QUEUE_ALIGN, 0);
	if (q->vaddr == NULL) {
		device_printf(sc->dev, "failed to allocate %d bytes\n", sz);
		return (-1);
	}

	q->mask = (1 << q->size_log2) - 1;
	q->head_off = (uint32_t)base - RISCV_IOMMU_CQB + RISCV_IOMMU_CQH;
	q->tail_off = (uint32_t)base - RISCV_IOMMU_CQB + RISCV_IOMMU_CQT;
	q->paddr = vtophys(q->vaddr);
	q->base = (sc->cq.size_log2 - 1) << CQB_LOG2SZ_1_S;
	q->base |= PHYS_TO_PPN(q->paddr) << CQB_PPN_S;
	WR8(sc, base, q->base);

	/* Verify it sticks. */
	reg = RD8(sc, base);
	if (reg != q->base) {
		device_printf(sc->dev, "could not init queue\n");
		return (ENXIO);
	}

	return (0);
}

static int
riscv_iommu_init_queues(struct riscv_iommu_softc *sc)
{
	int error;

	sc->cq.size_log2 = ilog2(CQ_ENTRY_COUNT);
	sc->fq.size_log2 = ilog2(FQ_ENTRY_COUNT);
	sc->pq.size_log2 = ilog2(PQ_ENTRY_COUNT);

	sc->cq.csr = RISCV_IOMMU_CQCSR;
	sc->fq.csr = RISCV_IOMMU_FQCSR;
	sc->pq.csr = RISCV_IOMMU_PQCSR;

	sc->cq.idx = 0;
	sc->fq.idx = 1;
	sc->pq.idx = 3;

	/* Command queue (CQ). */
	error = riscv_iommu_init_queue(sc, &sc->cq, RISCV_IOMMU_CQB,
	    CQ_ENTRY_DWORDS);
	if (error)
		return (error);

	/* Fault queue (FQ). */
	error = riscv_iommu_init_queue(sc, &sc->fq, RISCV_IOMMU_FQB,
	    FQ_ENTRY_DWORDS);
	if (error)
		return (error);

	/* Page request queue (PQ). */
	error = riscv_iommu_init_queue(sc, &sc->pq, RISCV_IOMMU_PQB,
	    PQ_ENTRY_DWORDS);
	if (error)
		return (error);

	error = riscv_iommu_enable_queue(sc, &sc->cq);
	if (error)
		return (error);

	error = riscv_iommu_enable_queue(sc, &sc->fq);
	if (error)
		return (error);

	error = riscv_iommu_enable_queue(sc, &sc->pq);
	if (error)
		return (error);

	return (0);
}

static int
riscv_iommu_init_pagedir(struct riscv_iommu_softc *sc)
{

	return (0);
}

static void
riscv_iommu_print_fault(struct riscv_iommu_softc *sc,
    struct riscv_iommu_fq_record *rec)
{
	struct riscv_iommu_fq_event *ev;
	uint16_t cause_id;
	uint16_t ttyp;
	uint32_t did;
	uint32_t pid;
	bool pv, priv;
	int i;

	cause_id = (rec->hdr & FQR_HDR_CAUSE_M) >> FQR_HDR_CAUSE_S;
	ttyp = (rec->hdr & FQR_HDR_TTYP_M) >> FQR_HDR_TTYP_S;
	did = (rec->hdr & FQR_HDR_DID_M) >> FQR_HDR_DID_S;
	pid = (rec->hdr & FQR_HDR_PID_M) >> FQR_HDR_PID_S;
	pv = (rec->hdr & FQR_HDR_PV) ? 1 : 0;
	priv = (rec->hdr & FQR_HDR_PRIV) ? 1 : 0;

	ev = NULL;
	for (i = 0; fq_events[i].cause_id != 0; i++) {
		if (fq_events[i].cause_id == cause_id) {
			ev = &fq_events[i];
			break;
		}
	}

	if (ev == NULL) {
		device_printf(sc->dev, "Fault: unknown fault 0x%x received\n",
		    cause_id);
		return;
	}

	device_printf(sc->dev, "Fault: event 0x%x received: %s\n",
	    ev->cause_id, ev->descr);
	device_printf(sc->dev, "    hdr 0x%lx\n", rec->hdr);
	device_printf(sc->dev, "    iotval 0x%lx\n", rec->iotval);
	device_printf(sc->dev, "    iotval2 0x%lx\n", rec->iotval2);
	device_printf(sc->dev, "    ttyp 0x%x did 0x%x pid 0x%x pv %d priv %d"
	    "\n", ttyp, did, pid, pv, priv);
}

static int
riscv_cq_intr(void *arg)
{
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_queue *q;
	uint32_t reg;

	sc = arg;
	q = &sc->cq;

	reg = RD4(sc, q->csr);
	printf("%s: pending %x\n", __func__, reg);

	/* Clear pending bit. */
	WR4(sc, RISCV_IOMMU_IPSR, IPSR_CIP);

	return (FILTER_HANDLED);
}

static int
riscv_fq_intr(void *arg)
{
	struct riscv_iommu_fq_record rec;
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_queue *q;
	uint32_t reg;

	sc = arg;
	q = &sc->fq;

	reg = RD4(sc, q->csr);
	printf("%s: pending %x\n", __func__, reg);

	/* Clear pending bit. */
	WR4(sc, RISCV_IOMMU_IPSR, IPSR_FIP);

	do {
		riscv_iommu_dequeue(sc, q, &rec);
		riscv_iommu_print_fault(sc, &rec);
	} while (!riscv_iommu_q_empty(q));

	return (FILTER_HANDLED);
}

static int
riscv_pm_intr(void *arg)
{
	struct riscv_iommu_softc *sc;

	sc = arg;

	printf("%s\n", __func__);

	/* Clear pending bit. */
	WR4(sc, RISCV_IOMMU_IPSR, IPSR_PMIP);

	return (FILTER_HANDLED);
}

static int
riscv_pq_intr(void *arg)
{
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_queue *q;
	uint32_t reg;

	sc = arg;
	q = &sc->pq;

	reg = RD4(sc, q->csr);
	printf("%s: pending %x\n", __func__, reg);

	/* Clear pending bit. */
	WR4(sc, RISCV_IOMMU_IPSR, IPSR_PIP);

	return (FILTER_HANDLED);
}

static int
riscv_iommu_init_ddt_linear(struct riscv_iommu_softc *sc)
{
	struct riscv_iommu_ddt *ddt;
	uint64_t size;
	uint64_t reg;

	ddt = &sc->ddt;
	ddt->num_top_entries = (1 << sc->l0_did_bits);

	size = ddt->num_top_entries * (sc->dc_dwords << 3);

	if (bootverbose)
		device_printf(sc->dev, "linear ddt size %ld, num_top_entries "
		    "%d\n", size, ddt->num_top_entries);

	ddt->vaddr = contigmalloc(size, M_IOMMU,
	    M_WAITOK | M_ZERO,	/* flags */
	    0,			/* low */
	    (1ul << 48) - 1,	/* high */
	    size,		/* alignment */
	    0);			/* boundary */
	if (ddt->vaddr == NULL) {
		device_printf(sc->dev, "failed to allocate ddt\n");
		return (ENXIO);
	}

	reg = vtophys(ddt->vaddr);
	if (bootverbose)
		device_printf(sc->dev, "ddt base %p size %lx\n", ddt->vaddr,
		    size);
	ddt->base = PHYS_TO_PPN(reg) << DDTP_PPN_S;

	return (0);
}

static int
riscv_iommu_init_ddt_2lvl(struct riscv_iommu_softc *sc)
{
	struct riscv_iommu_ddt *ddt;
	uint64_t size;
	uint64_t reg;
	uint64_t sz;

	ddt = &sc->ddt;
	ddt->num_top_entries = (1 << DDT_L1_DID_BITS);

	size = ddt->num_top_entries * (DDT_NON_LEAF_DWORDS << 3);

	if (bootverbose)
		device_printf(sc->dev, "%s: size %lu, l1 entries %d, size "
		    "%lu\n", __func__, size, ddt->num_top_entries, size);

	ddt->vaddr = contigmalloc(size, M_IOMMU,
	    M_WAITOK | M_ZERO,	/* flags */
	    0,			/* low */
	    (1ul << 48) - 1,	/* high */
	    size,		/* alignment */
	    0);			/* boundary */
	if (ddt->vaddr == NULL) {
		device_printf(sc->dev, "Failed to allocate 2lvl ddt.\n");
		return (ENOMEM);
	}

	sz = ddt->num_top_entries * sizeof(struct l1_desc);
	ddt->l1 = malloc(sz, M_IOMMU, M_WAITOK | M_ZERO);

	reg = vtophys(ddt->vaddr);
	if (bootverbose)
		device_printf(sc->dev, "ddt base %p size %lx\n", ddt->vaddr,
		    size);
	ddt->base = PHYS_TO_PPN(reg) << DDTP_PPN_S;

	return (0);
}

static int
riscv_iommu_init_l0_directory(struct riscv_iommu_softc *sc, int sid)
{
	struct riscv_iommu_ddt *ddt;
	struct l1_desc *l1_desc;
	uint64_t *l1e;
	uint64_t val;
	size_t size;
	int i;

	ddt = &sc->ddt;
	l1_desc = &ddt->l1[sid >> sc->l0_did_bits];
	if (l1_desc->va) {
		/* Already allocated. */
		return (0);
	}

	size = (1 << sc->l0_did_bits) * (sc->dc_dwords << 3);

	l1_desc->va = contigmalloc(size, M_IOMMU,
	    M_WAITOK | M_ZERO,	/* flags */
	    0,			/* low */
	    (1ul << 48) - 1,	/* high */
	    size,		/* alignment */
	    0);			/* boundary */
	if (l1_desc->va == NULL) {
		device_printf(sc->dev, "failed to allocate l0 directory\n");
		return (ENXIO);
	}

	l1_desc->pa = vtophys(l1_desc->va);

	i = sid >> sc->l0_did_bits;
	l1e = (void *)((uint64_t)ddt->vaddr + DDT_NON_LEAF_DWORDS * 8 * i);

	/* Install the L1 entry. */
	val = PHYS_TO_PPN(l1_desc->pa) << DC_NON_LEAF_ENTRY_PPN_S;
	val |= DC_NON_LEAF_ENTRY_VALID;
	*l1e = val;

	return (0);
}

static void *
riscv_iommu_get_dc_addr(struct riscv_iommu_softc *sc, int did)
{
	struct riscv_iommu_ddt *ddt;
	struct l1_desc *l1_desc;
	uintptr_t l0_base;
	void *addr;
	int l0_offs;
	int l1_idx;

	ddt = &sc->ddt;

	l0_offs = sc->dc_dwords * 8 * (did & ((1 << sc->l0_did_bits) - 1));

	if (sc->iommu_mode == DDTP_IOMMU_MODE_2LVL) {
		l1_idx = (did >> sc->l0_did_bits) &
		    ((1 << DDT_L1_DID_BITS) - 1);
		l1_desc = &ddt->l1[l1_idx];
		l0_base = (uintptr_t)l1_desc->va;
	} else
		l0_base = (uintptr_t)ddt->vaddr;

	addr = (void *)(l0_base + l0_offs);

	dprintf("ddt vaddr %p addr %p\n", ddt->vaddr, addr);

	return (addr);
}

static int
riscv_iommu_init_dc(struct riscv_iommu_softc *sc,
    struct riscv_iommu_domain *domain, int did, bool bypass)
{
	struct riscv_iommu_dc_base *dc_base;
	struct riscv_iommu_dc *dc;
	struct riscv_iommu_pmap *p;

	dc = riscv_iommu_get_dc_addr(sc, did);
	dc_base = &dc->base;

	device_printf(sc->dev, "address translation for device id"
	    " 0x%x is %s.\n", did, bypass ? "bypassed" : "enabled");

	p = &domain->p;

	bzero(dc_base, sizeof(struct riscv_iommu_dc_base));
	if (bypass == false)
		dc_base->fsc = p->pm_satp;
	dc_base->ta = (domain->pscid << DC_TA_PSCID_S) | DC_TA_V;

	riscv_iommu_inval_ddt_did(sc, did);
	riscv_iommu_sync(sc, &sc->cq);
	dc_base->tc |= DC_TC_V;
	riscv_iommu_inval_ddt_did(sc, did);
	riscv_iommu_inval_vma(sc);
	riscv_iommu_sync(sc, &sc->cq);

	return (0);
}

static void
riscv_iommu_deinit_dc(struct riscv_iommu_softc *sc, int did)
{
	struct riscv_iommu_dc_base *dc_base;
	struct riscv_iommu_dc *dc;

	dc = riscv_iommu_get_dc_addr(sc, did);
	dc_base = &dc->base;
	dc_base->tc &= ~DC_TC_V;

	riscv_iommu_inval_ddt_did(sc, did);
	riscv_iommu_sync(sc, &sc->cq);
}

static int
riscv_iommu_setup_interrupts(struct riscv_iommu_softc *sc)
{
	device_t dev;
	int error;

	dev = sc->dev;

	if (sc->res[1] == NULL || sc->res[2] == NULL ||
	    sc->res[3] == NULL || sc->res[4] == NULL) {
		device_printf(dev, "Warning: no interrupt resources "
		    "provided.\n");
		return (ENXIO);
	}

	error = bus_setup_intr(dev, sc->res[1], INTR_TYPE_MISC,
	    riscv_cq_intr, NULL, sc, &sc->intr_cookie[0]);
	if (error) {
		device_printf(dev, "Couldn't setup cq interrupt handler\n");
		return (ENXIO);
	}

	error = bus_setup_intr(dev, sc->res[2], INTR_TYPE_MISC,
	    riscv_fq_intr, NULL, sc, &sc->intr_cookie[1]);
	if (error) {
		device_printf(dev, "Couldn't setup fq interrupt handler\n");
		return (ENXIO);
	}

	error = bus_setup_intr(dev, sc->res[3], INTR_TYPE_MISC,
	    riscv_pm_intr, NULL, sc, &sc->intr_cookie[2]);
	if (error) {
		device_printf(dev, "Couldn't setup pm interrupt handler\n");
		return (ENXIO);
	}

	error = bus_setup_intr(dev, sc->res[4], INTR_TYPE_MISC,
	    riscv_pq_intr, NULL, sc, &sc->intr_cookie[3]);
	if (error) {
		device_printf(dev, "Couldn't setup pq interrupt handler\n");
		return (ENXIO);
	}

	WR8(sc, RISCV_IOMMU_ICVEC, 0 << 0 | 1 << 4 | 2 << 8 | 3 << 12);

	return (0);
}

int
riscv_iommu_attach(device_t dev)
{
	struct riscv_iommu_softc *sc;
	uint64_t caps;
	int error;

	sc = device_get_softc(dev);

	caps = bus_read_8(sc->res[0], RISCV_IOMMU_CAPABILITIES);
	if (bootverbose)
		device_printf(sc->dev, "IOMMU Capabilities: %lx\n", caps);

	device_printf(sc->dev, "Device-Context structure is %s.\n",
	    caps & CAPABILITIES_MSI_FLAT ?
	    "64-bytes (ext format)" : "32-bytes (std format)");

	if (caps & CAPABILITIES_MSI_FLAT) {
		sc->dc_dwords = DDT_DC_EXT_DWORDS;
		sc->l0_did_bits = 6;
	} else {
		sc->dc_dwords = DDT_DC_STD_DWORDS;
		sc->l0_did_bits = 7;
	}

	if (caps & CAPABILITIES_SV48)
		sc->pm_mode = PMAP_MODE_SV48;
	else if (caps & CAPABILITIES_SV39)
		sc->pm_mode = PMAP_MODE_SV39;
	else {
		device_printf(sc->dev, "Unsupported virtual memory system\n");
		return (ENXIO);
	}

	mtx_init(&sc->mtx, device_get_nameunit(sc->dev), "riscv_iommu",
	    MTX_DEF);

	WR4(sc, RISCV_IOMMU_FCTL, FCTL_WSI);

	error = riscv_iommu_setup_interrupts(sc);
	if (error) {
		device_printf(sc->dev, "Could not setup interrupts. "
		    "Continuing with no interrupts support.");
	}

	error = riscv_iommu_init_pagedir(sc);
	if (error)
		return (error);

	error = riscv_iommu_init_queues(sc);
	if (error)
		return (error);

	sc->iommu_mode = DDTP_IOMMU_MODE_2LVL;

	switch (sc->iommu_mode) {
	case DDTP_IOMMU_MODE_1LVL:
		error = riscv_iommu_init_ddt_linear(sc);
		break;
	case DDTP_IOMMU_MODE_2LVL:
		error = riscv_iommu_init_ddt_2lvl(sc);
		break;
	default:
		error = ENXIO;
	}
	if (error)
		return (error);

	sc->pscid_bits = 8;

	riscv_iommu_init_pscids(sc);
	if (error)
		return (error);

	error = riscv_iommu_set_mode(sc);
	if (error)
		return (error);

	return (0);
}

static int
riscv_iommu_set_buswide(device_t dev, struct riscv_iommu_domain *domain,
    struct riscv_iommu_ctx *ctx)
{
	struct riscv_iommu_softc *sc;
	int i;

	sc = device_get_softc(dev);

	printf("%s\n", __func__);

	for (i = 0; i < PCI_SLOTMAX; i++)
		riscv_iommu_init_dc(sc, domain, (ctx->did | i),
		    ctx->bypass);

	return (0);
}

static int
riscv_iommu_pci_get_did(device_t child, uintptr_t *xref0, u_int *did0)
{
	struct pci_id_ofw_iommu pi;
	int err;

	dprintf("%s\n", __func__);

	err = pci_get_id(child, PCI_ID_OFW_IOMMU, (uintptr_t *)&pi);
	if (err == 0) {
		if (did0)
			*did0 = pi.id;
		if (xref0)
			*xref0 = pi.xref;
	}

	return (err);
}

static int
riscv_iommu_find(device_t dev, device_t child)
{
	struct riscv_iommu_softc *sc;
	uintptr_t xref;
	int err;

	dprintf("%s\n", __func__);

	sc = device_get_softc(dev);

	err = riscv_iommu_pci_get_did(child, &xref, NULL);
	if (err)
		return (ENOENT);

	/* Check if xref is ours. */
	dprintf("xref %lx sc->xref %lx\n", xref, sc->xref);
	if (xref != sc->xref)
		return (EFAULT);

	return (0);
}

struct riscv_iommu_ctx *
riscv_iommu_ctx_lookup_by_did(device_t dev, u_int did)
{
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_unit *unit;
	struct riscv_iommu_ctx *ctx;

	dprintf("%s\n", __func__);
	sc = device_get_softc(dev);

	unit = &sc->unit;

	LIST_FOREACH(domain, &unit->domain_list, next) {
		LIST_FOREACH(ctx, &domain->ctx_list, next) {
			if (ctx->did == did) {
				refcount_acquire(&ctx->refcnt);
				return (ctx);
			}
		}
	}

	return (NULL);
}

static struct iommu_ctx *
riscv_iommu_ctx_lookup(device_t dev, device_t child)
{
	struct iommu_unit *iommu __diagused;
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_unit *unit;
	struct riscv_iommu_ctx *ctx;

	dprintf("%s\n", __func__);
	sc = device_get_softc(dev);

	unit = &sc->unit;
	iommu = &unit->iommu;

	IOMMU_ASSERT_LOCKED(iommu);

	LIST_FOREACH(domain, &unit->domain_list, next) {
		IOMMU_DOMAIN_LOCK(&domain->iodom);
		LIST_FOREACH(ctx, &domain->ctx_list, next) {
			if (ctx->dev == child) {
				refcount_acquire(&ctx->refcnt);
				IOMMU_DOMAIN_UNLOCK(&domain->iodom);
				return (&ctx->ioctx);
			}
		}
		IOMMU_DOMAIN_UNLOCK(&domain->iodom);
	}

	return (NULL);
}

static int
riscv_iommu_unmap(device_t dev, struct iommu_domain *iodom,
    vm_offset_t va, bus_size_t size)
{
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_softc *sc;
	int err;
	int i;

	sc = device_get_softc(dev);

	domain = (struct riscv_iommu_domain *)iodom;

	err = 0;

	dprintf("%s: %lx, %ld, domain %d\n", __func__, va, size, domain->pscid);

	for (i = 0; i < size; i += PAGE_SIZE) {
		if (iommu_pmap_remove(&domain->p, va) == 0) {
			/* pmap entry removed, invalidate TLB. */
			riscv_iommu_inval_vma_page(sc, va, domain->pscid);
		} else {
			err = ENOENT;
			break;
		}
		va += PAGE_SIZE;
	}

	riscv_iommu_sync(sc, &sc->cq);

	return (err);
}

static int
riscv_iommu_map(device_t dev, struct iommu_domain *iodom,
    vm_offset_t va, vm_page_t *ma, vm_size_t size,
    vm_prot_t prot)
{
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_softc *sc;
	vm_paddr_t pa;
	int error;
	int i;

	sc = device_get_softc(dev);

	domain = (struct riscv_iommu_domain *)iodom;

	for (i = 0; size > 0; size -= PAGE_SIZE) {
		pa = VM_PAGE_TO_PHYS(ma[i++]);
		dprintf("%s: %lx -> %lx, %ld, domain %d\n", __func__, va, pa,
		    size, domain->pscid);
		error = iommu_pmap_enter(&domain->p, va, pa, prot, 0);
		if (error)
			return (error);
		riscv_iommu_inval_vma_page(sc, va, domain->pscid);
		va += PAGE_SIZE;
	}

	riscv_iommu_sync(sc, &sc->cq);

	return (0);
}

static struct iommu_domain *
riscv_iommu_domain_alloc(device_t dev, struct iommu_unit *iommu)
{
	struct iommu_domain *iodom;
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_unit *unit;
	struct riscv_iommu_softc *sc;
	int new_pscid;
	int va_bits;
	int error;

	sc = device_get_softc(dev);

	dprintf("%s\n", __func__);

	unit = (struct riscv_iommu_unit *)iommu;

	error = riscv_iommu_pscid_alloc(sc, &new_pscid);
	if (error) {
		device_printf(sc->dev,
		   "Could not allocate PSCID for a new domain.\n");
		return (NULL);
	}

	domain = malloc(sizeof(*domain), M_IOMMU, M_WAITOK | M_ZERO);
	domain->pscid = (uint16_t)new_pscid;

	iommu_pmap_pinit(&domain->p, sc->pm_mode);

	riscv_iommu_inval_vma_pscid(sc, domain->pscid);

	LIST_INIT(&domain->ctx_list);

	IOMMU_LOCK(iommu);
	LIST_INSERT_HEAD(&unit->domain_list, domain, next);
	IOMMU_UNLOCK(iommu);

	iodom = &domain->iodom;

	va_bits = sc->pm_mode == PMAP_MODE_SV48 ? 48 : 39;

	/* Avoid sign-extension. */
	va_bits -= 1;

	iodom->end = (1ULL << va_bits) - 1;

	return (iodom);
}

static void
riscv_iommu_domain_free(device_t dev, struct iommu_domain *iodom)
{
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_softc *sc;

	sc = device_get_softc(dev);

	dprintf("%s\n", __func__);

	domain = (struct riscv_iommu_domain *)iodom;

	LIST_REMOVE(domain, next);

	iommu_pmap_remove_pages(&domain->p);
	iommu_pmap_release(&domain->p);

	riscv_iommu_inval_vma_pscid(sc, domain->pscid);
	riscv_iommu_pscid_free(sc, domain->pscid);

	free(domain, M_IOMMU);
}

static struct iommu_ctx *
riscv_iommu_ctx_alloc(device_t dev, struct iommu_domain *iodom, device_t child,
    bool disabled)
{
	struct riscv_iommu_domain *domain;
	struct riscv_iommu_ctx *ctx;

	dprintf("%s\n", __func__);

	domain = (struct riscv_iommu_domain *)iodom;

	ctx = malloc(sizeof(struct riscv_iommu_ctx), M_IOMMU,
	    M_WAITOK | M_ZERO);
	ctx->dev = child;
	ctx->domain = domain;
	refcount_init(&ctx->refcnt, 1);
	if (disabled)
		ctx->bypass = true;

	IOMMU_DOMAIN_LOCK(iodom);
	LIST_INSERT_HEAD(&domain->ctx_list, ctx, next);
	IOMMU_DOMAIN_UNLOCK(iodom);

	return (&ctx->ioctx);
}

static int
riscv_iommu_ctx_init(device_t dev, struct iommu_ctx *ioctx)
{
	struct riscv_iommu_domain *domain;
	struct iommu_domain *iodom;
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_ctx *ctx;
	devclass_t pci_class;
	u_int did;
	int error;

	ctx = (struct riscv_iommu_ctx *)ioctx;

	dprintf("%s\n", __func__);

	sc = device_get_softc(dev);

	domain = ctx->domain;
	iodom = (struct iommu_domain *)domain;

	pci_class = devclass_find("pci");
	if (device_get_devclass(device_get_parent(ctx->dev)) == pci_class) {
		error = riscv_iommu_pci_get_did(ctx->dev, NULL, &did);
		if (error)
			return (error);

		ioctx->rid = pci_get_rid(dev);
		ctx->did = did;
		ctx->vendor = pci_get_vendor(ctx->dev);
		ctx->device = pci_get_device(ctx->dev);
	}

	if (sc->iommu_mode == DDTP_IOMMU_MODE_2LVL) {
		error = riscv_iommu_init_l0_directory(sc, ctx->did);
		if (error)
			return (error);
	}
	riscv_iommu_init_dc(sc, domain, ctx->did, ctx->bypass);

	if (device_get_devclass(device_get_parent(ctx->dev)) == pci_class)
		if (iommu_is_buswide_ctx(iodom->iommu, pci_get_bus(ctx->dev)))
			riscv_iommu_set_buswide(dev, domain, ctx);

	return (0);
}

static bool
riscv_iommu_ctx_free(device_t dev, struct iommu_ctx *ioctx)
{
	struct riscv_iommu_softc *sc;
	struct riscv_iommu_ctx *ctx;

	dprintf("%s\n", __func__);

	IOMMU_ASSERT_LOCKED(ioctx->domain->iommu);

	sc = device_get_softc(dev);

	ctx = (struct riscv_iommu_ctx *)ioctx;
	if (refcount_release(&ctx->refcnt)) {
		riscv_iommu_deinit_dc(sc, ctx->did);
		LIST_REMOVE(ctx, next);
		free(ctx, M_IOMMU);
		return (true);
	}

	return (false);
}

#ifdef FDT
static int
riscv_iommu_ofw_md_data(device_t dev, struct iommu_ctx *ioctx, pcell_t *cells,
    int ncells)
{
	struct riscv_iommu_ctx *ctx;

	printf("%s\n", __func__);
	ctx = (struct riscv_iommu_ctx *)ioctx;

	if (ncells != 1)
		return (-1);

	ctx->did = cells[0];

	return (0);
}
#endif

static int
riscv_iommu_read_ivar(device_t dev, device_t child, int which,
    uintptr_t *result)
{
	struct riscv_iommu_softc *sc;

	sc = device_get_softc(dev);

	device_printf(sc->dev, "%s\n", __func__);

	return (ENOENT);
}

static device_method_t riscv_iommu_methods[] = {
	/* IOMMU interface */
	DEVMETHOD(iommu_find,		riscv_iommu_find),
	DEVMETHOD(iommu_map,		riscv_iommu_map),
	DEVMETHOD(iommu_unmap,		riscv_iommu_unmap),
	DEVMETHOD(iommu_domain_alloc,	riscv_iommu_domain_alloc),
	DEVMETHOD(iommu_domain_free,	riscv_iommu_domain_free),
	DEVMETHOD(iommu_ctx_alloc,	riscv_iommu_ctx_alloc),
	DEVMETHOD(iommu_ctx_init,	riscv_iommu_ctx_init),
	DEVMETHOD(iommu_ctx_free,	riscv_iommu_ctx_free),
	DEVMETHOD(iommu_ctx_lookup,	riscv_iommu_ctx_lookup),
#ifdef FDT
	DEVMETHOD(iommu_ofw_md_data,	riscv_iommu_ofw_md_data),
#endif

	/* Bus interface */
	DEVMETHOD(bus_read_ivar,	riscv_iommu_read_ivar),

	/* End */
	DEVMETHOD_END
};

DEFINE_CLASS_0(riscv_iommu, riscv_iommu_driver, riscv_iommu_methods,
    sizeof(struct riscv_iommu_softc));