// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2019 Western Digital Corporation or its affiliates.
 * Copyright (c) 2025 Ventana Micro Systems Inc.
 */

#include <linux/bitops.h>
#include <linux/cpufeature.h>
#include <linux/errno.h>
#include <linux/kvm_host.h>
#include <linux/module.h>
#include <linux/pgtable.h>
#include <asm/kvm_gstage.h>
#include <asm/hwcap.h>

#ifdef CONFIG_64BIT
unsigned long kvm_riscv_gstage_max_pgd_levels __ro_after_init = 3;
#else
unsigned long kvm_riscv_gstage_max_pgd_levels __ro_after_init = 2;
#endif

#define gstage_pte_leaf(__ptep)	\
	(pte_val(*(__ptep)) & (_PAGE_READ | _PAGE_WRITE | _PAGE_EXEC))

static inline unsigned long gstage_pte_index(struct kvm_gstage *gstage,
					     gpa_t addr, u32 level)
{
	unsigned long mask;
	unsigned long shift = HGATP_PAGE_SHIFT + (kvm_riscv_gstage_index_bits * level);

	if (level == gstage->pgd_levels - 1)
		mask = (PTRS_PER_PTE * (1UL << kvm_riscv_gstage_pgd_xbits)) - 1;
	else
		mask = PTRS_PER_PTE - 1;

	return (addr >> shift) & mask;
}

static inline unsigned long gstage_pte_page_vaddr(pte_t pte)
{
	return (unsigned long)pfn_to_virt(__page_val_to_pfn(pte_val(pte)));
}

static int gstage_page_size_to_level(struct kvm_gstage *gstage, unsigned long page_size,
				     u32 *out_level)
{
	u32 i;
	unsigned long psz = 1UL << 12;

	for (i = 0; i < gstage->pgd_levels; i++) {
		if (page_size == (psz << (i * kvm_riscv_gstage_index_bits))) {
			*out_level = i;
			return 0;
		}
	}

	return -EINVAL;
}

static int gstage_level_to_page_order(struct kvm_gstage *gstage, u32 level,
				      unsigned long *out_pgorder)
{
	if (gstage->pgd_levels < level)
		return -EINVAL;

	*out_pgorder = 12 + (level * kvm_riscv_gstage_index_bits);
	return 0;
}

static int gstage_level_to_page_size(struct kvm_gstage *gstage, u32 level,
				     unsigned long *out_pgsize)
{
	int rc;
	unsigned long page_order = PAGE_SHIFT;

	rc = gstage_level_to_page_order(gstage, level, &page_order);
	if (rc)
		return rc;

	*out_pgsize = BIT(page_order);
	return 0;
}

bool kvm_riscv_gstage_get_leaf(struct kvm_gstage *gstage, gpa_t addr,
			       pte_t **ptepp, u32 *ptep_level)
{
	pte_t *ptep;
	u32 current_level = gstage->pgd_levels - 1;

	*ptep_level = current_level;
	ptep = (pte_t *)gstage->pgd;
	ptep = &ptep[gstage_pte_index(gstage, addr, current_level)];
	while (ptep && pte_val(ptep_get(ptep))) {
		if (gstage_pte_leaf(ptep)) {
			*ptep_level = current_level;
			*ptepp = ptep;
			return true;
		}

		if (current_level) {
			current_level--;
			*ptep_level = current_level;
			ptep = (pte_t *)gstage_pte_page_vaddr(ptep_get(ptep));
			ptep = &ptep[gstage_pte_index(gstage, addr, current_level)];
		} else {
			ptep = NULL;
		}
	}

	return false;
}

static void gstage_tlb_flush(struct kvm_gstage *gstage, u32 level, gpa_t addr)
{
	unsigned long order = PAGE_SHIFT;

	if (gstage_level_to_page_order(gstage, level, &order))
		return;
	addr &= ~(BIT(order) - 1);

	if (gstage->flags & KVM_GSTAGE_FLAGS_LOCAL)
		kvm_riscv_local_hfence_gvma_vmid_gpa(gstage->vmid, addr, BIT(order), order);
	else
		kvm_riscv_hfence_gvma_vmid_gpa(gstage->kvm, -1UL, 0, addr, BIT(order), order,
					       gstage->vmid);
}

bool kvm_riscv_gstage_try_update_pte(struct kvm_gstage *gstage, u32 level,
				     gpa_t addr, pte_t *ptep,
				     pte_t old_pte, pte_t new_pte)
{
	if (cmpxchg(&ptep->pte, pte_val(old_pte), pte_val(new_pte)) !=
	    pte_val(old_pte))
		return false;

	if (pte_val(old_pte) != pte_val(new_pte))
		gstage_tlb_flush(gstage, level, addr);

	return true;
}

int kvm_riscv_gstage_set_pte(struct kvm_gstage *gstage,
			     struct kvm_mmu_memory_cache *pcache,
			     const struct kvm_gstage_mapping *map)
{
	u32 current_level = gstage->pgd_levels - 1;
	pte_t *next_ptep = (pte_t *)gstage->pgd;
	pte_t *ptep = &next_ptep[gstage_pte_index(gstage, map->addr, current_level)];

	if (current_level < map->level)
		return -EINVAL;

	while (current_level != map->level) {
		if (gstage_pte_leaf(ptep))
			return -EEXIST;

		if (!pte_val(ptep_get(ptep))) {
			if (!pcache)
				return -ENOMEM;
			next_ptep = kvm_mmu_memory_cache_alloc(pcache);
			if (!next_ptep)
				return -ENOMEM;
			set_pte(ptep, pfn_pte(PFN_DOWN(__pa(next_ptep)),
					      __pgprot(_PAGE_TABLE)));
		} else {
			if (gstage_pte_leaf(ptep))
				return -EEXIST;
			next_ptep = (pte_t *)gstage_pte_page_vaddr(ptep_get(ptep));
		}

		current_level--;
		ptep = &next_ptep[gstage_pte_index(gstage, map->addr, current_level)];
	}

	if (pte_val(*ptep) != pte_val(map->pte)) {
		bool was_invalid = !pte_val(*ptep);
		set_pte(ptep, map->pte);
		if (gstage_pte_leaf(ptep) &&
		    !(was_invalid && riscv_has_extension_unlikely(RISCV_ISA_EXT_SVVPTC)))
			gstage_tlb_flush(gstage, current_level, map->addr);
	}

	return 0;
}

static void kvm_riscv_gstage_update_pte_prot(struct kvm_gstage *gstage, u32 level,
					     gpa_t addr, pte_t *ptep, pgprot_t prot)
{
	pte_t old_pte, new_pte;

	for (;;) {
		old_pte = ptep_get(ptep);
		if (pgprot_val(pte_pgprot(old_pte)) == pgprot_val(prot))
			return;

		new_pte = pfn_pte(pte_pfn(old_pte), prot);
		new_pte = pte_mkdirty(new_pte);

		if (kvm_riscv_gstage_try_update_pte(gstage, level, addr, ptep,
						    old_pte, new_pte))
			return;

		cpu_relax();
	}
}

int kvm_riscv_gstage_map_page(struct kvm_gstage *gstage,
			      struct kvm_mmu_memory_cache *pcache,
			      gpa_t gpa, phys_addr_t hpa, unsigned long page_size,
			      bool page_rdonly, bool page_exec,
			      struct kvm_gstage_mapping *out_map)
{
	bool found_leaf;
	u32 ptep_level;
	pgprot_t prot;
	pte_t *ptep;
	int ret;

	out_map->addr = gpa;
	out_map->level = 0;

	ret = gstage_page_size_to_level(gstage, page_size, &out_map->level);
	if (ret)
		return ret;

	/*
	 * A RISC-V implementation can choose to either:
	 * 1) Update 'A' and 'D' PTE bits in hardware
	 * 2) Generate page fault when 'A' and/or 'D' bits are not set
	 *    PTE so that software can update these bits.
	 *
	 * We support both options mentioned above. To achieve this, we
	 * always set 'A' and 'D' PTE bits at time of creating G-stage
	 * mapping. To support KVM dirty page logging with both options
	 * mentioned above, we will write-protect G-stage PTEs to track
	 * dirty pages.
	 */

	if (page_exec) {
		if (page_rdonly)
			prot = PAGE_READ_EXEC;
		else
			prot = PAGE_WRITE_EXEC;
	} else {
		if (page_rdonly)
			prot = PAGE_READ;
		else
			prot = PAGE_WRITE;
	}

	found_leaf = kvm_riscv_gstage_get_leaf(gstage, gpa, &ptep, &ptep_level);
	if (found_leaf) {
		/*
		 * ptep_level is the current gstage mapping level of addr, out_map->level
		 * is the required mapping level during fault handling.
		 *
		 * 1) ptep_level > out_map->level
		 * This happens when dirty logging is enabled and huge pages are used.
		 * KVM must track the pages at 4K level, and split the huge mapping
		 * into 4K mappings.
		 *
		 * 2) ptep_level < out_map->level
		 * This happens when dirty logging is disabled and huge pages are used.
		 * The gstage is split into 4K mappings, but the out_map level is now
		 * back to the huge page level. Ignore the out_map level this time, and
		 * just update the pte prot here. Otherwise, we would fall back to mapping
		 * the gstage at huge page level in `kvm_riscv_gstage_set_pte`, with the
		 * overhead of freeing the page tables(not support now), which would slow
		 * down the vCPUs' performance.
		 *
		 * It is better to recover the huge page mapping in the ioctl context when
		 * disabling dirty logging.
		 *
		 * 3) ptep_level == out_map->level
		 * We already have the ptep, just update the pte prot if the pfn not change.
		 * There is no need to invoke `kvm_riscv_gstage_set_pte` again.
		 */
		if (ptep_level > out_map->level) {
			kvm_riscv_gstage_split_huge(gstage, pcache, gpa,
						    out_map->level, true);
		} else if (ALIGN_DOWN(PFN_PHYS(pte_pfn(ptep_get(ptep))), page_size) == hpa) {
			kvm_riscv_gstage_update_pte_prot(gstage, ptep_level, gpa, ptep, prot);
			return 0;
		}
	}

	out_map->pte = pfn_pte(PFN_DOWN(hpa), prot);
	out_map->pte = pte_mkdirty(out_map->pte);

	return kvm_riscv_gstage_set_pte(gstage, pcache, out_map);
}

static inline unsigned long make_child_pte(unsigned long huge_pte, int index,
					   unsigned long child_page_size)
{
	unsigned long child_pte = huge_pte;
	unsigned long child_pfn_offset;

	/*
	 * The child_pte already has the base address of the huge page being
	 * split. So we just have to OR in the offset to the page at the next
	 * lower level for the given index.
	 */
	child_pfn_offset = index * (child_page_size / PAGE_SIZE);
	child_pte |= pte_val(pfn_pte(child_pfn_offset, __pgprot(0)));

	return child_pte;
}

int kvm_riscv_gstage_split_huge(struct kvm_gstage *gstage,
				struct kvm_mmu_memory_cache *pcache,
				gpa_t addr, u32 target_level, bool flush)
{
	u32 current_level = gstage->pgd_levels - 1;
	pte_t *next_ptep = (pte_t *)gstage->pgd;
	unsigned long huge_pte, child_pte;
	unsigned long child_page_size;
	pte_t *ptep;
	int i, ret;

	if (!pcache)
		return -ENOMEM;

	while(current_level > target_level) {
		ptep = (pte_t *)&next_ptep[gstage_pte_index(gstage, addr, current_level)];

		if (!pte_val(ptep_get(ptep)))
			break;

		if (!gstage_pte_leaf(ptep)) {
			next_ptep = (pte_t *)gstage_pte_page_vaddr(ptep_get(ptep));
			current_level--;
			continue;
		}

		huge_pte = pte_val(ptep_get(ptep));

		ret = gstage_level_to_page_size(gstage, current_level - 1, &child_page_size);
		if (ret)
			return ret;

		next_ptep = kvm_mmu_memory_cache_alloc(pcache);
		if (!next_ptep)
			return -ENOMEM;

		for (i = 0; i < PTRS_PER_PTE; i++) {
			child_pte = make_child_pte(huge_pte, i, child_page_size);
			set_pte((pte_t *)&next_ptep[i], __pte(child_pte));
		}

		set_pte(ptep, pfn_pte(PFN_DOWN(__pa(next_ptep)),
				__pgprot(_PAGE_TABLE)));

		if (flush)
			gstage_tlb_flush(gstage, current_level, addr);

		current_level--;
	}

	return 0;
}

bool kvm_riscv_gstage_op_pte(struct kvm_gstage *gstage, gpa_t addr,
			     pte_t *ptep, u32 ptep_level, enum kvm_riscv_gstage_op op)
{
	int i, ret;
	pte_t old_pte, *next_ptep;
	u32 next_ptep_level;
	unsigned long next_page_size, page_size;
	bool flush = false;

	ret = gstage_level_to_page_size(gstage, ptep_level, &page_size);
	if (ret)
		return false;

	WARN_ON(addr & (page_size - 1));

	if (!pte_val(ptep_get(ptep)))
		return false;

	if (ptep_level && !gstage_pte_leaf(ptep)) {
		next_ptep = (pte_t *)gstage_pte_page_vaddr(ptep_get(ptep));
		next_ptep_level = ptep_level - 1;
		ret = gstage_level_to_page_size(gstage, next_ptep_level, &next_page_size);
		if (ret)
			return false;

		if (op == GSTAGE_OP_CLEAR)
			set_pte(ptep, __pte(0));
		for (i = 0; i < PTRS_PER_PTE; i++)
			flush |= kvm_riscv_gstage_op_pte(gstage, addr + i * next_page_size,
							 &next_ptep[i], next_ptep_level, op);
		if (op == GSTAGE_OP_CLEAR)
			put_page(virt_to_page(next_ptep));
	} else {
		old_pte = *ptep;
		if (op == GSTAGE_OP_CLEAR)
			set_pte(ptep, __pte(0));
		else if (op == GSTAGE_OP_WP)
			set_pte(ptep, __pte(pte_val(ptep_get(ptep)) & ~_PAGE_WRITE));
		if (pte_val(*ptep) != pte_val(old_pte))
			flush = true;
	}

	return flush;
}

bool kvm_riscv_gstage_unmap_range(struct kvm_gstage *gstage,
				  gpa_t start, gpa_t size, bool may_block)
{
	int ret;
	pte_t *ptep;
	u32 ptep_level;
	bool found_leaf;
	unsigned long page_size;
	gpa_t addr = start, end = start + size;
	bool flush = false;

	while (addr < end) {
		found_leaf = kvm_riscv_gstage_get_leaf(gstage, addr, &ptep, &ptep_level);
		ret = gstage_level_to_page_size(gstage, ptep_level, &page_size);
		if (ret)
			break;

		if (!found_leaf) {
			addr = ALIGN(addr + 1, page_size);
		} else {
			if (!(addr & (page_size - 1)) && ((end - addr) >= page_size))
				flush |= kvm_riscv_gstage_op_pte(gstage, addr, ptep,
								 ptep_level, GSTAGE_OP_CLEAR);
			else {
				WARN_ONCE(1, "Skip unmap range addr: %#llx, end: %#llx, page_size: %#lx\n",
						addr, end, page_size);
			}

			addr += page_size;
		}

		/*
		 * If the range is too large, release the kvm->mmu_lock
		 * to prevent starvation and lockup detector warnings.
		 */
		if (!(gstage->flags & KVM_GSTAGE_FLAGS_LOCAL) && may_block && addr < end)
			cond_resched_rwlock_write(&gstage->kvm->mmu_lock);
	}

	return flush;
}

bool kvm_riscv_gstage_wp_range(struct kvm_gstage *gstage, gpa_t start, gpa_t end)
{
	int ret;
	pte_t *ptep;
	u32 ptep_level;
	bool found_leaf;
	gpa_t addr = start;
	unsigned long page_size;
	bool flush = false;

	while (addr < end) {
		found_leaf = kvm_riscv_gstage_get_leaf(gstage, addr, &ptep, &ptep_level);
		ret = gstage_level_to_page_size(gstage, ptep_level, &page_size);
		if (ret)
			break;

		if (!found_leaf) {
			addr = ALIGN(addr + 1, page_size);
		} else {
			addr = ALIGN_DOWN(addr, page_size);
			flush |= kvm_riscv_gstage_op_pte(gstage, addr, ptep,
							 ptep_level, GSTAGE_OP_WP);
			addr += page_size;
		}
	}

	return flush;
}

void __init kvm_riscv_gstage_mode_detect(void)
{
#ifdef CONFIG_64BIT
	/* Try Sv57x4 G-stage mode */
	csr_write(CSR_HGATP, HGATP_MODE_SV57X4 << HGATP_MODE_SHIFT);
	if ((csr_read(CSR_HGATP) >> HGATP_MODE_SHIFT) == HGATP_MODE_SV57X4) {
		kvm_riscv_gstage_max_pgd_levels = 5;
		goto done;
	}

	/* Try Sv48x4 G-stage mode */
	csr_write(CSR_HGATP, HGATP_MODE_SV48X4 << HGATP_MODE_SHIFT);
	if ((csr_read(CSR_HGATP) >> HGATP_MODE_SHIFT) == HGATP_MODE_SV48X4) {
		kvm_riscv_gstage_max_pgd_levels = 4;
		goto done;
	}

	/* Try Sv39x4 G-stage mode */
	csr_write(CSR_HGATP, HGATP_MODE_SV39X4 << HGATP_MODE_SHIFT);
	if ((csr_read(CSR_HGATP) >> HGATP_MODE_SHIFT) == HGATP_MODE_SV39X4) {
		kvm_riscv_gstage_max_pgd_levels = 3;
		goto done;
	}
#else /* CONFIG_32BIT */
	/* Try Sv32x4 G-stage mode */
	csr_write(CSR_HGATP, HGATP_MODE_SV32X4 << HGATP_MODE_SHIFT);
	if ((csr_read(CSR_HGATP) >> HGATP_MODE_SHIFT) == HGATP_MODE_SV32X4) {
		kvm_riscv_gstage_max_pgd_levels = 2;
		goto done;
	}
#endif

	/* KVM depends on !HGATP_MODE_OFF */
	kvm_riscv_gstage_max_pgd_levels = 0;

done:
	csr_write(CSR_HGATP, 0);
	kvm_riscv_local_hfence_gvma_all();
}