// SPDX-License-Identifier: GPL-2.0
/*
 *  KVM guest address space mapping code
 *
 *    Copyright IBM Corp. 2007, 2020, 2024
 *    Author(s): Claudio Imbrenda <imbrenda@linux.ibm.com>
 *		 Martin Schwidefsky <schwidefsky@de.ibm.com>
 *		 David Hildenbrand <david@redhat.com>
 *		 Janosch Frank <frankja@linux.ibm.com>
 */

#include <linux/kernel.h>
#include <linux/pagewalk.h>
#include <linux/swap.h>
#include <linux/smp.h>
#include <linux/spinlock.h>
#include <linux/slab.h>
#include <linux/swapops.h>
#include <linux/ksm.h>
#include <linux/mm.h>
#include <linux/mman.h>
#include <linux/pgtable.h>
#include <linux/kvm_types.h>
#include <linux/kvm_host.h>
#include <linux/pgalloc.h>

#include <asm/page-states.h>
#include <asm/tlb.h>
#include "dat.h"

int kvm_s390_mmu_cache_topup(struct kvm_s390_mmu_cache *mc)
{
	void *o;

	for ( ; mc->n_crsts < KVM_S390_MMU_CACHE_N_CRSTS; mc->n_crsts++) {
		o = (void *)__get_free_pages(GFP_KERNEL_ACCOUNT | __GFP_COMP, CRST_ALLOC_ORDER);
		if (!o)
			return -ENOMEM;
		mc->crsts[mc->n_crsts] = o;
	}
	for ( ; mc->n_pts < KVM_S390_MMU_CACHE_N_PTS; mc->n_pts++) {
		o = (void *)__get_free_page(GFP_KERNEL_ACCOUNT);
		if (!o)
			return -ENOMEM;
		mc->pts[mc->n_pts] = o;
	}
	for ( ; mc->n_rmaps < KVM_S390_MMU_CACHE_N_RMAPS; mc->n_rmaps++) {
		o = kzalloc_obj(struct vsie_rmap, GFP_KERNEL_ACCOUNT);
		if (!o)
			return -ENOMEM;
		mc->rmaps[mc->n_rmaps] = o;
	}
	return 0;
}

static inline struct page_table *dat_alloc_pt_noinit(struct kvm_s390_mmu_cache *mc)
{
	struct page_table *res;

	res = kvm_s390_mmu_cache_alloc_pt(mc);
	if (res)
		__arch_set_page_dat(res, 1);
	return res;
}

static inline struct crst_table *dat_alloc_crst_noinit(struct kvm_s390_mmu_cache *mc)
{
	struct crst_table *res;

	res = kvm_s390_mmu_cache_alloc_crst(mc);
	if (res)
		__arch_set_page_dat(res, 1UL << CRST_ALLOC_ORDER);
	return res;
}

struct crst_table *dat_alloc_crst_sleepable(unsigned long init)
{
	struct page *page;
	void *virt;

	page = alloc_pages(GFP_KERNEL_ACCOUNT | __GFP_COMP, CRST_ALLOC_ORDER);
	if (!page)
		return NULL;
	virt = page_to_virt(page);
	__arch_set_page_dat(virt, 1UL << CRST_ALLOC_ORDER);
	crst_table_init(virt, init);
	return virt;
}

void dat_free_level(struct crst_table *table, bool owns_ptes)
{
	unsigned int i;

	for (i = 0; i < _CRST_ENTRIES; i++) {
		if (table->crstes[i].h.fc || table->crstes[i].h.i)
			continue;
		if (!is_pmd(table->crstes[i]))
			dat_free_level(dereference_crste(table->crstes[i]), owns_ptes);
		else if (owns_ptes)
			dat_free_pt(dereference_pmd(table->crstes[i].pmd));
	}
	dat_free_crst(table);
}

int dat_set_asce_limit(struct kvm_s390_mmu_cache *mc, union asce *asce, int newtype)
{
	struct crst_table *table;
	union crste crste;

	while (asce->dt > newtype) {
		table = dereference_asce(*asce);
		crste = table->crstes[0];
		if (crste.h.fc)
			return 0;
		if (!crste.h.i) {
			asce->rsto = crste.h.fc0.to;
			dat_free_crst(table);
		} else {
			crste.h.tt--;
			crst_table_init((void *)table, crste.val);
		}
		asce->dt--;
	}
	while (asce->dt < newtype) {
		crste = _crste_fc0(asce->rsto, asce->dt + 1);
		table = dat_alloc_crst_noinit(mc);
		if (!table)
			return -ENOMEM;
		crst_table_init((void *)table, _CRSTE_HOLE(crste.h.tt).val);
		table->crstes[0] = crste;
		asce->rsto = __pa(table) >> PAGE_SHIFT;
		asce->dt++;
	}
	return 0;
}

/**
 * dat_crstep_xchg_atomic() - Atomically exchange a gmap CRSTE with another.
 * @crstep: Pointer to the CRST entry.
 * @old: Expected old value.
 * @new: Replacement entry.
 * @gfn: The affected guest address.
 * @asce: The asce of the address space.
 *
 * This function is needed to atomically exchange a CRSTE that potentially
 * maps a prefix area, without having to invalidate it inbetween.
 *
 * Context: This function is assumed to be called with kvm->mmu_lock held.
 *
 * Return: %true if the exchange was successful.
 */
bool __must_check dat_crstep_xchg_atomic(union crste *crstep, union crste old, union crste new,
					 gfn_t gfn, union asce asce)
{
	if (old.h.i)
		return arch_try_cmpxchg((long *)crstep, &old.val, new.val);
	if (cpu_has_edat2())
		return crdte_crste(crstep, old, new, gfn, asce);
	return cspg_crste(crstep, old, new);
}

static void dat_set_storage_key_from_pgste(union pte pte, union pgste pgste)
{
	union skey nkey = { .acc = pgste.acc, .fp = pgste.fp };

	page_set_storage_key(pte_origin(pte), nkey.skey, 0);
}

static void dat_move_storage_key(union pte old, union pte new)
{
	page_set_storage_key(pte_origin(new), page_get_storage_key(pte_origin(old)), 1);
}

static union pgste dat_save_storage_key_into_pgste(union pte pte, union pgste pgste)
{
	union skey skey;

	skey.skey = page_get_storage_key(pte_origin(pte));

	pgste.acc = skey.acc;
	pgste.fp = skey.fp;
	pgste.gr |= skey.r;
	pgste.gc |= skey.c;

	return pgste;
}

union pgste __dat_ptep_xchg(union pte *ptep, union pgste pgste, union pte new, gfn_t gfn,
			    union asce asce, bool uses_skeys)
{
	union pte old = READ_ONCE(*ptep);

	/* Updating only the software bits while holding the pgste lock. */
	if (!((ptep->val ^ new.val) & ~_PAGE_SW_BITS)) {
		WRITE_ONCE(ptep->swbyte, new.swbyte);
		return pgste;
	}

	if (!old.h.i) {
		unsigned long opts = IPTE_GUEST_ASCE | (pgste.nodat ? IPTE_NODAT : 0);

		if (machine_has_tlb_guest())
			__ptep_ipte(gfn_to_gpa(gfn), (void *)ptep, opts, asce.val, IPTE_GLOBAL);
		else
			__ptep_ipte(gfn_to_gpa(gfn), (void *)ptep, 0, 0, IPTE_GLOBAL);
	}

	if (uses_skeys) {
		if (old.h.i && !new.h.i)
			/* Invalid to valid: restore storage keys from PGSTE. */
			dat_set_storage_key_from_pgste(new, pgste);
		else if (!old.h.i && new.h.i)
			/* Valid to invalid: save storage keys to PGSTE. */
			pgste = dat_save_storage_key_into_pgste(old, pgste);
		else if (!old.h.i && !new.h.i)
			/* Valid to valid: move storage keys. */
			if (old.h.pfra != new.h.pfra)
				dat_move_storage_key(old, new);
		/* Invalid to invalid: nothing to do. */
	}

	WRITE_ONCE(*ptep, new);
	return pgste;
}

/*
 * dat_split_ste() - Split a segment table entry into page table entries.
 *
 * Context: This function is assumed to be called with kvm->mmu_lock held.
 *
 * Return: 0 in case of success, -ENOMEM if running out of memory.
 */
static int dat_split_ste(struct kvm_s390_mmu_cache *mc, union pmd *pmdp, gfn_t gfn,
			 union asce asce, bool uses_skeys)
{
	union pgste pgste_init;
	struct page_table *pt;
	union pmd new, old;
	union pte init;
	int i;

	BUG_ON(!mc);
	old = READ_ONCE(*pmdp);

	/* Already split, nothing to do. */
	if (!old.h.i && !old.h.fc)
		return 0;

	pt = dat_alloc_pt_noinit(mc);
	if (!pt)
		return -ENOMEM;
	new.val = virt_to_phys(pt);

	while (old.h.i || old.h.fc) {
		init.val = pmd_origin_large(old);
		init.h.p = old.h.p;
		init.h.i = old.h.i;
		init.s.d = old.s.fc1.d;
		init.s.w = old.s.fc1.w;
		init.s.y = old.s.fc1.y;
		init.s.sd = old.s.fc1.sd;
		init.s.pr = old.s.fc1.pr;
		pgste_init.val = 0;
		if (old.h.fc) {
			for (i = 0; i < _PAGE_ENTRIES; i++)
				pt->ptes[i].val = init.val | i * PAGE_SIZE;
			/* No need to take locks as the page table is not installed yet. */
			pgste_init.prefix_notif = old.s.fc1.prefix_notif;
			pgste_init.vsie_notif = old.s.fc1.vsie_notif;
			pgste_init.vsie_gmem = old.s.fc1.vsie_notif;
			pgste_init.pcl = uses_skeys && init.h.i;
			dat_init_pgstes(pt, pgste_init.val);
		} else {
			dat_init_page_table(pt, init.val, 0);
		}

		if (dat_pmdp_xchg_atomic(pmdp, old, new, gfn, asce)) {
			if (!pgste_init.pcl)
				return 0;
			for (i = 0; i < _PAGE_ENTRIES; i++) {
				union pgste pgste = pt->pgstes[i];

				pgste = dat_save_storage_key_into_pgste(pt->ptes[i], pgste);
				pgste_set_unlock(pt->ptes + i, pgste);
			}
			return 0;
		}
		old = READ_ONCE(*pmdp);
	}

	dat_free_pt(pt);
	return 0;
}

/*
 * dat_split_crste() - Split a crste into smaller crstes.
 *
 * Context: This function is assumed to be called with kvm->mmu_lock held.
 *
 * Return: %0 in case of success, %-ENOMEM if running out of memory.
 */
static int dat_split_crste(struct kvm_s390_mmu_cache *mc, union crste *crstep,
			   gfn_t gfn, union asce asce, bool uses_skeys)
{
	struct crst_table *table;
	union crste old, new, init;
	int i;

	old = READ_ONCE(*crstep);
	if (is_pmd(old))
		return dat_split_ste(mc, &crstep->pmd, gfn, asce, uses_skeys);

	BUG_ON(!mc);

	/* Already split, nothing to do. */
	if (!old.h.i && !old.h.fc)
		return 0;

	table = dat_alloc_crst_noinit(mc);
	if (!table)
		return -ENOMEM;

	new.val = virt_to_phys(table);
	new.h.tt = old.h.tt;
	new.h.fc0.tl = _REGION_ENTRY_LENGTH;

	while (old.h.i || old.h.fc) {
		init = old;
		init.h.tt--;
		if (old.h.fc) {
			for (i = 0; i < _CRST_ENTRIES; i++)
				table->crstes[i].val = init.val | i * HPAGE_SIZE;
		} else {
			crst_table_init((void *)table, init.val);
		}
		if (dat_crstep_xchg_atomic(crstep, old, new, gfn, asce))
			return 0;
		old = READ_ONCE(*crstep);
	}

	dat_free_crst(table);
	return 0;
}

/**
 * dat_entry_walk() - Walk the gmap page tables.
 * @mc: Cache to use to allocate dat tables, if needed; can be NULL if neither
 *      %DAT_WALK_SPLIT or %DAT_WALK_ALLOC is specified in @flags.
 * @gfn: Guest frame.
 * @asce: The ASCE of the address space.
 * @flags: Flags from WALK_* macros.
 * @walk_level: Level to walk to, from LEVEL_* macros.
 * @last: Will be filled the last visited non-pte DAT entry.
 * @ptepp: Will be filled the last visited pte entry, if any, otherwise NULL.
 *
 * Returns a table entry pointer for the given guest address and @walk_level.
 *
 * The @flags have the following meanings:
 * * %DAT_WALK_IGN_HOLES: consider holes as normal table entries
 * * %DAT_WALK_ALLOC: allocate new tables to reach the requested level, if needed
 * * %DAT_WALK_SPLIT: split existing large pages to reach the requested level, if needed
 * * %DAT_WALK_LEAF: return successfully whenever a large page is encountered
 * * %DAT_WALK_ANY: return successfully even if the requested level could not be reached
 * * %DAT_WALK_CONTINUE: walk to the requested level with the specified flags, and then try to
 *                       continue walking to ptes with only DAT_WALK_ANY
 * * %DAT_WALK_USES_SKEYS: storage keys are in use
 *
 * Context: called with kvm->mmu_lock held.
 *
 * Return:
 * * %PGM_ADDRESSING if the requested address lies outside memory
 * * a PIC number if the requested address lies in a memory hole of type _DAT_TOKEN_PIC
 * * %-EFAULT if the requested address lies inside a memory hole of a different type
 * * %-EINVAL if the given ASCE is not compatible with the requested level
 * * %-EFBIG if the requested level could not be reached because a larger frame was found
 * * %-ENOENT if the requested level could not be reached for other reasons
 * * %-ENOMEM if running out of memory while allocating or splitting a table
 */
int dat_entry_walk(struct kvm_s390_mmu_cache *mc, gfn_t gfn, union asce asce, int flags,
		   int walk_level, union crste **last, union pte **ptepp)
{
	union vaddress vaddr = { .addr = gfn_to_gpa(gfn) };
	bool continue_anyway = flags & DAT_WALK_CONTINUE;
	bool uses_skeys = flags & DAT_WALK_USES_SKEYS;
	bool ign_holes = flags & DAT_WALK_IGN_HOLES;
	bool allocate = flags & DAT_WALK_ALLOC;
	bool split = flags & DAT_WALK_SPLIT;
	bool leaf = flags & DAT_WALK_LEAF;
	bool any = flags & DAT_WALK_ANY;
	struct page_table *pgtable;
	struct crst_table *table;
	union crste entry;
	int rc;

	*last = NULL;
	*ptepp = NULL;
	if (WARN_ON_ONCE(unlikely(!asce.val)))
		return -EINVAL;
	if (WARN_ON_ONCE(unlikely(walk_level > asce.dt)))
		return -EINVAL;
	if (!asce_contains_gfn(asce, gfn))
		return PGM_ADDRESSING;

	table = dereference_asce(asce);
	if (asce.dt >= ASCE_TYPE_REGION1) {
		*last = table->crstes + vaddr.rfx;
		entry = READ_ONCE(**last);
		if (WARN_ON_ONCE(entry.h.tt != TABLE_TYPE_REGION1))
			return -EINVAL;
		if (crste_hole(entry) && !ign_holes)
			return entry.tok.type == _DAT_TOKEN_PIC ? entry.tok.par : -EFAULT;
		if (walk_level == TABLE_TYPE_REGION1)
			return 0;
		if (entry.pgd.h.i) {
			if (!allocate)
				return any ? 0 : -ENOENT;
			rc = dat_split_crste(mc, *last, gfn, asce, uses_skeys);
			if (rc)
				return rc;
			entry = READ_ONCE(**last);
		}
		table = dereference_crste(entry.pgd);
	}

	if (asce.dt >= ASCE_TYPE_REGION2) {
		*last = table->crstes + vaddr.rsx;
		entry = READ_ONCE(**last);
		if (WARN_ON_ONCE(entry.h.tt != TABLE_TYPE_REGION2))
			return -EINVAL;
		if (crste_hole(entry) && !ign_holes)
			return entry.tok.type == _DAT_TOKEN_PIC ? entry.tok.par : -EFAULT;
		if (walk_level == TABLE_TYPE_REGION2)
			return 0;
		if (entry.p4d.h.i) {
			if (!allocate)
				return any ? 0 : -ENOENT;
			rc = dat_split_crste(mc, *last, gfn, asce, uses_skeys);
			if (rc)
				return rc;
			entry = READ_ONCE(**last);
		}
		table = dereference_crste(entry.p4d);
	}

	if (asce.dt >= ASCE_TYPE_REGION3) {
		*last = table->crstes + vaddr.rtx;
		entry = READ_ONCE(**last);
		if (WARN_ON_ONCE(entry.h.tt != TABLE_TYPE_REGION3))
			return -EINVAL;
		if (crste_hole(entry) && !ign_holes)
			return entry.tok.type == _DAT_TOKEN_PIC ? entry.tok.par : -EFAULT;
		if (walk_level == TABLE_TYPE_REGION3 &&
		    continue_anyway && !entry.pud.h.fc && !entry.h.i) {
			walk_level = TABLE_TYPE_PAGE_TABLE;
			allocate = false;
		}
		if (walk_level == TABLE_TYPE_REGION3 || ((leaf || any) && entry.pud.h.fc))
			return 0;
		if (entry.pud.h.i && !entry.pud.h.fc) {
			if (!allocate)
				return any ? 0 : -ENOENT;
			rc = dat_split_crste(mc, *last, gfn, asce, uses_skeys);
			if (rc)
				return rc;
			entry = READ_ONCE(**last);
		}
		if (walk_level <= TABLE_TYPE_SEGMENT && entry.pud.h.fc) {
			if (!split)
				return -EFBIG;
			rc = dat_split_crste(mc, *last, gfn, asce, uses_skeys);
			if (rc)
				return rc;
			entry = READ_ONCE(**last);
		}
		table = dereference_crste(entry.pud);
	}

	*last = table->crstes + vaddr.sx;
	entry = READ_ONCE(**last);
	if (WARN_ON_ONCE(entry.h.tt != TABLE_TYPE_SEGMENT))
		return -EINVAL;
	if (crste_hole(entry) && !ign_holes)
		return entry.tok.type == _DAT_TOKEN_PIC ? entry.tok.par : -EFAULT;
	if (continue_anyway && !entry.pmd.h.fc && !entry.h.i) {
		walk_level = TABLE_TYPE_PAGE_TABLE;
		allocate = false;
	}
	if (walk_level == TABLE_TYPE_SEGMENT || ((leaf || any) && entry.pmd.h.fc))
		return 0;

	if (entry.pmd.h.i && !entry.pmd.h.fc) {
		if (!allocate)
			return any ? 0 : -ENOENT;
		rc = dat_split_ste(mc, &(*last)->pmd, gfn, asce, uses_skeys);
		if (rc)
			return rc;
		entry = READ_ONCE(**last);
	}
	if (walk_level <= TABLE_TYPE_PAGE_TABLE && entry.pmd.h.fc) {
		if (!split)
			return -EFBIG;
		rc = dat_split_ste(mc, &(*last)->pmd, gfn, asce, uses_skeys);
		if (rc)
			return rc;
		entry = READ_ONCE(**last);
	}
	pgtable = dereference_pmd(entry.pmd);
	*ptepp = pgtable->ptes + vaddr.px;
	if (pte_hole(**ptepp) && !ign_holes)
		return (*ptepp)->tok.type == _DAT_TOKEN_PIC ? (*ptepp)->tok.par : -EFAULT;
	return 0;
}

static long dat_pte_walk_range(gfn_t gfn, gfn_t end, struct page_table *table, struct dat_walk *w)
{
	unsigned int idx = gfn & (_PAGE_ENTRIES - 1);
	long rc = 0;

	for ( ; gfn < end; idx++, gfn++) {
		if (pte_hole(READ_ONCE(table->ptes[idx]))) {
			if (!(w->flags & DAT_WALK_IGN_HOLES))
				return -EFAULT;
			if (!(w->flags & DAT_WALK_ANY))
				continue;
		}

		rc = w->ops->pte_entry(table->ptes + idx, gfn, gfn + 1, w);
		if (rc)
			break;
	}
	return rc;
}

static long dat_crste_walk_range(gfn_t start, gfn_t end, struct crst_table *table,
				 struct dat_walk *walk)
{
	unsigned long idx, cur_shift, cur_size;
	dat_walk_op the_op;
	union crste crste;
	gfn_t cur, next;
	long rc = 0;

	cur_shift = 8 + table->crstes[0].h.tt * 11;
	idx = (start >> cur_shift) & (_CRST_ENTRIES - 1);
	cur_size = 1UL << cur_shift;

	for (cur = ALIGN_DOWN(start, cur_size); cur < end; idx++, cur = next) {
		next = cur + cur_size;
		walk->last = table->crstes + idx;
		crste = READ_ONCE(*walk->last);

		if (crste_hole(crste)) {
			if (!(walk->flags & DAT_WALK_IGN_HOLES))
				return -EFAULT;
			if (!(walk->flags & DAT_WALK_ANY))
				continue;
		}

		the_op = walk->ops->crste_ops[crste.h.tt];
		if (the_op) {
			rc = the_op(walk->last, cur, next, walk);
			crste = READ_ONCE(*walk->last);
		}
		if (rc)
			break;
		if (!crste.h.i && !crste.h.fc) {
			if (!is_pmd(crste))
				rc = dat_crste_walk_range(max(start, cur), min(end, next),
							  _dereference_crste(crste), walk);
			else if (walk->ops->pte_entry)
				rc = dat_pte_walk_range(max(start, cur), min(end, next),
							dereference_pmd(crste.pmd), walk);
			if (rc)
				break;
		}
	}
	return rc;
}

/**
 * _dat_walk_gfn_range() - Walk DAT tables.
 * @start: The first guest page frame to walk.
 * @end: The guest page frame immediately after the last one to walk.
 * @asce: The ASCE of the guest mapping.
 * @ops: The gmap_walk_ops that will be used to perform the walk.
 * @flags: Flags from WALK_* (currently only WALK_IGN_HOLES is supported).
 * @priv: Will be passed as-is to the callbacks.
 *
 * Any callback returning non-zero causes the walk to stop immediately.
 *
 * Return: %-EINVAL in case of error, %-EFAULT if @start is too high for the
 *         given ASCE unless the DAT_WALK_IGN_HOLES flag is specified,
 *         otherwise it returns whatever the callbacks return.
 */
long _dat_walk_gfn_range(gfn_t start, gfn_t end, union asce asce,
			 const struct dat_walk_ops *ops, int flags, void *priv)
{
	struct crst_table *table = dereference_asce(asce);
	struct dat_walk walk = {
		.ops	= ops,
		.asce	= asce,
		.priv	= priv,
		.flags	= flags,
		.start	= start,
		.end	= end,
	};

	if (WARN_ON_ONCE(unlikely(!asce.val)))
		return -EINVAL;
	if (!asce_contains_gfn(asce, start))
		return (flags & DAT_WALK_IGN_HOLES) ? 0 : -EFAULT;

	return dat_crste_walk_range(start, min(end, asce_end(asce)), table, &walk);
}

int dat_get_storage_key(union asce asce, gfn_t gfn, union skey *skey)
{
	union crste *crstep;
	union pgste pgste;
	union pte *ptep;
	int rc;

	skey->skey = 0;
	rc = dat_entry_walk(NULL, gfn, asce, DAT_WALK_ANY, TABLE_TYPE_PAGE_TABLE, &crstep, &ptep);
	if (rc)
		return rc;

	if (!ptep) {
		union crste crste;

		crste = READ_ONCE(*crstep);
		if (!crste.h.fc || !crste.s.fc1.pr)
			return 0;
		skey->skey = page_get_storage_key(large_crste_to_phys(crste, gfn));
		return 0;
	}
	pgste = pgste_get_lock(ptep);
	if (ptep->h.i) {
		skey->acc = pgste.acc;
		skey->fp = pgste.fp;
	} else {
		skey->skey = page_get_storage_key(pte_origin(*ptep));
	}
	skey->r |= pgste.gr;
	skey->c |= pgste.gc;
	pgste_set_unlock(ptep, pgste);
	return 0;
}

static void dat_update_ptep_sd(union pgste old, union pgste pgste, union pte *ptep)
{
	if (pgste.acc != old.acc || pgste.fp != old.fp || pgste.gr != old.gr || pgste.gc != old.gc)
		__atomic64_or(_PAGE_SD, &ptep->val);
}

int dat_set_storage_key(struct kvm_s390_mmu_cache *mc, union asce asce, gfn_t gfn,
			union skey skey, bool nq)
{
	union pgste pgste, old;
	union crste *crstep;
	union pte *ptep;
	int rc;

	rc = dat_entry_walk(mc, gfn, asce, DAT_WALK_LEAF_ALLOC, TABLE_TYPE_PAGE_TABLE,
			    &crstep, &ptep);
	if (rc)
		return rc;

	if (!ptep) {
		page_set_storage_key(large_crste_to_phys(*crstep, gfn), skey.skey, !nq);
		return 0;
	}

	old = pgste_get_lock(ptep);
	pgste = old;

	pgste.acc = skey.acc;
	pgste.fp = skey.fp;
	pgste.gc = skey.c;
	pgste.gr = skey.r;

	if (!ptep->h.i) {
		union skey old_skey;

		old_skey.skey = page_get_storage_key(pte_origin(*ptep));
		pgste.hc |= old_skey.c;
		pgste.hr |= old_skey.r;
		old_skey.c = old.gc;
		old_skey.r = old.gr;
		skey.r = 0;
		skey.c = 0;
		page_set_storage_key(pte_origin(*ptep), skey.skey, !nq);
	}

	dat_update_ptep_sd(old, pgste, ptep);
	pgste_set_unlock(ptep, pgste);
	return 0;
}

static bool page_cond_set_storage_key(phys_addr_t paddr, union skey skey, union skey *oldkey,
				      bool nq, bool mr, bool mc)
{
	oldkey->skey = page_get_storage_key(paddr);
	if (oldkey->acc == skey.acc && oldkey->fp == skey.fp &&
	    (oldkey->r == skey.r || mr) && (oldkey->c == skey.c || mc))
		return false;
	page_set_storage_key(paddr, skey.skey, !nq);
	return true;
}

int dat_cond_set_storage_key(struct kvm_s390_mmu_cache *mmc, union asce asce, gfn_t gfn,
			     union skey skey, union skey *oldkey, bool nq, bool mr, bool mc)
{
	union pgste pgste, old;
	union crste *crstep;
	union skey prev;
	union pte *ptep;
	int rc;

	rc = dat_entry_walk(mmc, gfn, asce, DAT_WALK_LEAF_ALLOC, TABLE_TYPE_PAGE_TABLE,
			    &crstep, &ptep);
	if (rc)
		return rc;

	if (!ptep)
		return page_cond_set_storage_key(large_crste_to_phys(*crstep, gfn), skey, oldkey,
						 nq, mr, mc);

	old = pgste_get_lock(ptep);
	pgste = old;

	rc = 1;
	pgste.acc = skey.acc;
	pgste.fp = skey.fp;
	pgste.gc = skey.c;
	pgste.gr = skey.r;

	if (!ptep->h.i) {
		rc = page_cond_set_storage_key(pte_origin(*ptep), skey, &prev, nq, mr, mc);
		pgste.hc |= prev.c;
		pgste.hr |= prev.r;
		prev.c |= old.gc;
		prev.r |= old.gr;
	} else {
		prev.acc = old.acc;
		prev.fp = old.fp;
		prev.c = old.gc;
		prev.r = old.gr;
	}
	if (oldkey)
		*oldkey = prev;

	dat_update_ptep_sd(old, pgste, ptep);
	pgste_set_unlock(ptep, pgste);
	return rc;
}

int dat_reset_reference_bit(union asce asce, gfn_t gfn, union skey *skey)
{
	union pgste pgste, old;
	union crste *crstep;
	union pte *ptep;
	int rc;

	skey->skey = 0;

	rc = dat_entry_walk(NULL, gfn, asce, DAT_WALK_ANY, TABLE_TYPE_PAGE_TABLE, &crstep, &ptep);
	if (rc)
		return rc;

	if (!ptep) {
		union crste crste = READ_ONCE(*crstep);

		if (!crste.h.fc || !crste.s.fc1.pr)
			return 0;
		skey->skey = page_reset_referenced(large_crste_to_phys(*crstep, gfn)) << 1;
		return 0;
	}
	old = pgste_get_lock(ptep);
	pgste = old;

	if (!ptep->h.i) {
		skey->skey = page_reset_referenced(pte_origin(*ptep)) << 1;
		pgste.hr = skey->r;
	}
	skey->r |= pgste.gr;
	skey->c |= pgste.gc;
	pgste.gr = 0;

	dat_update_ptep_sd(old, pgste, ptep);
	pgste_set_unlock(ptep, pgste);
	return 0;
}

static long dat_reset_skeys_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	union pgste pgste;

	pgste = pgste_get_lock(ptep);
	pgste.acc = 0;
	pgste.fp = 0;
	pgste.gr = 0;
	pgste.gc = 0;
	if (ptep->s.pr)
		page_set_storage_key(pte_origin(*ptep), PAGE_DEFAULT_KEY, 1);
	pgste_set_unlock(ptep, pgste);

	if (need_resched())
		return next;
	return 0;
}

static long dat_reset_skeys_crste(union crste *crstep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	phys_addr_t addr, end, origin = crste_origin_large(*crstep);

	if (!crstep->h.fc || !crstep->s.fc1.pr)
		return 0;

	addr = ((max(gfn, walk->start) - gfn) << PAGE_SHIFT) + origin;
	end = ((min(next, walk->end) - gfn) << PAGE_SHIFT) + origin;
	while (ALIGN(addr + 1, _SEGMENT_SIZE) <= end)
		addr = sske_frame(addr, PAGE_DEFAULT_KEY);
	for ( ; addr < end; addr += PAGE_SIZE)
		page_set_storage_key(addr, PAGE_DEFAULT_KEY, 1);

	if (need_resched())
		return next;
	return 0;
}

long dat_reset_skeys(union asce asce, gfn_t start)
{
	const struct dat_walk_ops ops = {
		.pte_entry = dat_reset_skeys_pte,
		.pmd_entry = dat_reset_skeys_crste,
		.pud_entry = dat_reset_skeys_crste,
	};

	return _dat_walk_gfn_range(start, asce_end(asce), asce, &ops, DAT_WALK_IGN_HOLES, NULL);
}

struct slot_priv {
	unsigned long token;
	struct kvm_s390_mmu_cache *mc;
};

static long _dat_slot_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	struct slot_priv *p = walk->priv;
	union crste dummy = { .val = p->token };
	union pte new_pte, pte = READ_ONCE(*ptep);
	union pgste pgste;

	new_pte = _PTE_TOK(dummy.tok.type, dummy.tok.par);

	/* Table entry already in the desired state. */
	if (pte.val == new_pte.val)
		return 0;

	pgste = pgste_get_lock(ptep);
	pgste = __dat_ptep_xchg(ptep, pgste, new_pte, gfn, walk->asce, false);
	pgste.cmma_d = 0;
	pgste_set_unlock(ptep, pgste);

	return 0;
}

static long _dat_slot_crste(union crste *crstep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	union crste new_crste, crste = READ_ONCE(*crstep);
	struct slot_priv *p = walk->priv;

	new_crste.val = p->token;
	new_crste.h.tt = crste.h.tt;

	/* Table entry already in the desired state. */
	if (crste.val == new_crste.val)
		return 0;

	/* This table entry needs to be updated. */
	if (walk->start <= gfn && walk->end >= next) {
		if (!dat_crstep_xchg_atomic(crstep, crste, new_crste, gfn, walk->asce))
			return -EINVAL;
		/* A lower level table was present, needs to be freed. */
		if (!crste.h.fc && !crste.h.i) {
			if (is_pmd(crste))
				dat_free_pt(dereference_pmd(crste.pmd));
			else
				dat_free_level(dereference_crste(crste), true);
		}
		return 0;
	}

	/* A lower level table is present, things will handled there. */
	if (!crste.h.fc && !crste.h.i)
		return 0;
	/* Split (install a lower level table), and handle things there. */
	return dat_split_crste(p->mc, crstep, gfn, walk->asce, false);
}

static const struct dat_walk_ops dat_slot_ops = {
	.pte_entry = _dat_slot_pte,
	.crste_ops = { _dat_slot_crste, _dat_slot_crste, _dat_slot_crste, _dat_slot_crste, },
};

int dat_set_slot(struct kvm_s390_mmu_cache *mc, union asce asce, gfn_t start, gfn_t end,
		 u16 type, u16 param)
{
	struct slot_priv priv = {
		.token = _CRSTE_TOK(0, type, param).val,
		.mc = mc,
	};

	return _dat_walk_gfn_range(start, end, asce, &dat_slot_ops,
				   DAT_WALK_IGN_HOLES | DAT_WALK_ANY, &priv);
}

static void pgste_set_unlock_multiple(union pte *first, int n, union pgste *pgstes)
{
	int i;

	for (i = 0; i < n; i++) {
		if (!pgstes[i].pcl)
			break;
		pgste_set_unlock(first + i, pgstes[i]);
	}
}

static bool pgste_get_trylock_multiple(union pte *first, int n, union pgste *pgstes)
{
	int i;

	for (i = 0; i < n; i++) {
		if (!pgste_get_trylock(first + i, pgstes + i))
			break;
	}
	if (i == n)
		return true;
	pgste_set_unlock_multiple(first, n, pgstes);
	return false;
}

unsigned long dat_get_ptval(struct page_table *table, struct ptval_param param)
{
	union pgste pgstes[4] = {};
	unsigned long res = 0;
	int i, n;

	n = param.len + 1;

	while (!pgste_get_trylock_multiple(table->ptes + param.offset, n, pgstes))
		cpu_relax();

	for (i = 0; i < n; i++)
		res = res << 16 | pgstes[i].val16;

	pgste_set_unlock_multiple(table->ptes + param.offset, n, pgstes);
	return res;
}

void dat_set_ptval(struct page_table *table, struct ptval_param param, unsigned long val)
{
	union pgste pgstes[4] = {};
	int i, n;

	n = param.len + 1;

	while (!pgste_get_trylock_multiple(table->ptes + param.offset, n, pgstes))
		cpu_relax();

	for (i = param.len; i >= 0; i--) {
		pgstes[i].val16 = val;
		val = val >> 16;
	}

	pgste_set_unlock_multiple(table->ptes + param.offset, n, pgstes);
}

static long _dat_test_young_pte(union pte *ptep, gfn_t start, gfn_t end, struct dat_walk *walk)
{
	return ptep->s.y;
}

static long _dat_test_young_crste(union crste *crstep, gfn_t start, gfn_t end,
				  struct dat_walk *walk)
{
	return crstep->h.fc && crstep->s.fc1.y;
}

static const struct dat_walk_ops test_age_ops = {
	.pte_entry = _dat_test_young_pte,
	.pmd_entry = _dat_test_young_crste,
	.pud_entry = _dat_test_young_crste,
};

/**
 * dat_test_age_gfn() - Test young.
 * @asce: The ASCE whose address range is to be tested.
 * @start: The first guest frame of the range to check.
 * @end: The guest frame after the last in the range.
 *
 * Context: called by KVM common code with the kvm mmu write lock held.
 *
 * Return: %true if any page in the given range is young, otherwise %false.
 */
bool dat_test_age_gfn(union asce asce, gfn_t start, gfn_t end)
{
	return _dat_walk_gfn_range(start, end, asce, &test_age_ops, 0, NULL) > 0;
}

static long dat_set_pn_crste(union crste *crstep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	union crste newcrste, oldcrste;
	int *n = walk->priv;

	do {
		oldcrste = READ_ONCE(*crstep);
		if (!oldcrste.h.fc || oldcrste.h.i || oldcrste.h.p)
			return 0;
		if (oldcrste.s.fc1.prefix_notif)
			break;
		newcrste = oldcrste;
		newcrste.s.fc1.prefix_notif = 1;
	} while (!dat_crstep_xchg_atomic(crstep, oldcrste, newcrste, gfn, walk->asce));
	*n = 2;
	return 0;
}

static long dat_set_pn_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	int *n = walk->priv;
	union pgste pgste;

	pgste = pgste_get_lock(ptep);
	if (!ptep->h.i && !ptep->h.p) {
		pgste.prefix_notif = 1;
		*n += 1;
	}
	pgste_set_unlock(ptep, pgste);
	return 0;
}

int dat_set_prefix_notif_bit(union asce asce, gfn_t gfn)
{
	static const struct dat_walk_ops ops = {
		.pte_entry = dat_set_pn_pte,
		.pmd_entry = dat_set_pn_crste,
		.pud_entry = dat_set_pn_crste,
	};

	int n = 0;

	_dat_walk_gfn_range(gfn, gfn + 2, asce, &ops, DAT_WALK_IGN_HOLES, &n);
	if (n != 2)
		return -EAGAIN;
	return 0;
}

/**
 * dat_perform_essa() - Perform ESSA actions on the PGSTE.
 * @asce: The asce to operate on.
 * @gfn: The guest page frame to operate on.
 * @orc: The specific action to perform, see the ESSA_SET_* macros.
 * @state: The storage attributes to be returned to the guest.
 * @dirty: Returns whether the function dirtied a previously clean entry.
 *
 * Context: Called with kvm->mmu_lock held.
 *
 * Return:
 * * %1 if the page state has been altered and the page is to be added to the CBRL
 * * %0 if the page state has been altered, but the page is not to be added to the CBRL
 * * %-1 if the page state has not been altered and the page is not to be added to the CBRL
 */
int dat_perform_essa(union asce asce, gfn_t gfn, int orc, union essa_state *state, bool *dirty)
{
	union crste *crstep;
	union pgste pgste;
	union pte *ptep;
	int res = 0;

	if (dat_entry_walk(NULL, gfn, asce, 0, TABLE_TYPE_PAGE_TABLE, &crstep, &ptep)) {
		*state = (union essa_state) { .exception = 1 };
		return -1;
	}

	pgste = pgste_get_lock(ptep);

	*state = (union essa_state) {
		.content = (ptep->h.i << 1) + (ptep->h.i && pgste.zero),
		.nodat = pgste.nodat,
		.usage = pgste.usage,
		};

	switch (orc) {
	case ESSA_GET_STATE:
		res = -1;
		break;
	case ESSA_SET_STABLE:
		pgste.usage = PGSTE_GPS_USAGE_STABLE;
		pgste.nodat = 0;
		break;
	case ESSA_SET_UNUSED:
		pgste.usage = PGSTE_GPS_USAGE_UNUSED;
		if (ptep->h.i)
			res = 1;
		break;
	case ESSA_SET_VOLATILE:
		pgste.usage = PGSTE_GPS_USAGE_VOLATILE;
		if (ptep->h.i)
			res = 1;
		break;
	case ESSA_SET_POT_VOLATILE:
		if (!ptep->h.i) {
			pgste.usage = PGSTE_GPS_USAGE_POT_VOLATILE;
		} else if (pgste.zero) {
			pgste.usage = PGSTE_GPS_USAGE_VOLATILE;
		} else if (!pgste.gc) {
			pgste.usage = PGSTE_GPS_USAGE_VOLATILE;
			res = 1;
		}
		break;
	case ESSA_SET_STABLE_RESIDENT:
		pgste.usage = PGSTE_GPS_USAGE_STABLE;
		/*
		 * Since the resident state can go away any time after this
		 * call, we will not make this page resident. We can revisit
		 * this decision if a guest will ever start using this.
		 */
		break;
	case ESSA_SET_STABLE_IF_RESIDENT:
		if (!ptep->h.i)
			pgste.usage = PGSTE_GPS_USAGE_STABLE;
		break;
	case ESSA_SET_STABLE_NODAT:
		pgste.usage = PGSTE_GPS_USAGE_STABLE;
		pgste.nodat = 1;
		break;
	default:
		WARN_ONCE(1, "Invalid ORC!");
		res = -1;
		break;
	}
	/* If we are discarding a page, set it to logical zero. */
	pgste.zero = res == 1;
	if (orc > 0) {
		*dirty = !pgste.cmma_d;
		pgste.cmma_d = 1;
	}

	pgste_set_unlock(ptep, pgste);

	return res;
}

static long dat_reset_cmma_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	union pgste pgste;

	pgste = pgste_get_lock(ptep);
	pgste.usage = 0;
	pgste.nodat = 0;
	pgste.cmma_d = 0;
	pgste_set_unlock(ptep, pgste);
	if (need_resched())
		return next;
	return 0;
}

long dat_reset_cmma(union asce asce, gfn_t start)
{
	const struct dat_walk_ops dat_reset_cmma_ops = {
		.pte_entry = dat_reset_cmma_pte,
	};

	return _dat_walk_gfn_range(start, asce_end(asce), asce, &dat_reset_cmma_ops,
				   DAT_WALK_IGN_HOLES, NULL);
}

struct dat_get_cmma_state {
	gfn_t start;
	gfn_t end;
	unsigned int count;
	u8 *values;
	atomic64_t *remaining;
};

static long __dat_peek_cmma_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	struct dat_get_cmma_state *state = walk->priv;
	union pgste pgste;

	pgste = pgste_get_lock(ptep);
	state->values[gfn - walk->start] = pgste.usage | (pgste.nodat << 6);
	pgste_set_unlock(ptep, pgste);
	state->end = next;

	return 0;
}

static long __dat_peek_cmma_crste(union crste *crstep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	struct dat_get_cmma_state *state = walk->priv;

	if (crstep->h.i)
		state->end = min(walk->end, next);
	return 0;
}

int dat_peek_cmma(gfn_t start, union asce asce, unsigned int *count, u8 *values)
{
	const struct dat_walk_ops ops = {
		.pte_entry = __dat_peek_cmma_pte,
		.pmd_entry = __dat_peek_cmma_crste,
		.pud_entry = __dat_peek_cmma_crste,
		.p4d_entry = __dat_peek_cmma_crste,
		.pgd_entry = __dat_peek_cmma_crste,
	};
	struct dat_get_cmma_state state = { .values = values, };
	int rc;

	rc = _dat_walk_gfn_range(start, start + *count, asce, &ops, DAT_WALK_DEFAULT, &state);
	*count = state.end >= start ? state.end - start : 0;
	/* Return success if at least one value was saved, otherwise an error. */
	return (rc == -EFAULT && *count > 0) ? 0 : rc;
}

static long __dat_get_cmma_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	struct dat_get_cmma_state *state = walk->priv;
	union pgste pgste;

	if (state->start != -1) {
		if ((gfn - state->end) > KVM_S390_MAX_BIT_DISTANCE)
			return 1;
		if (gfn - state->start >= state->count)
			return 1;
	}

	if (!READ_ONCE(*pgste_of(ptep)).cmma_d)
		return 0;

	pgste = pgste_get_lock(ptep);
	if (pgste.cmma_d) {
		if (state->start == -1)
			state->start = gfn;
		pgste.cmma_d = 0;
		atomic64_dec(state->remaining);
		state->values[gfn - state->start] = pgste.usage | pgste.nodat << 6;
		state->end = next;
	}
	pgste_set_unlock(ptep, pgste);
	return 0;
}

int dat_get_cmma(union asce asce, gfn_t *start, unsigned int *count, u8 *values, atomic64_t *rem)
{
	const struct dat_walk_ops ops = { .pte_entry = __dat_get_cmma_pte, };
	struct dat_get_cmma_state state = {
		.remaining = rem,
		.values = values,
		.count = *count,
		.start = -1,
	};

	_dat_walk_gfn_range(*start, asce_end(asce), asce, &ops, DAT_WALK_IGN_HOLES, &state);
	/* If no dirty pages were found, wrap around and continue searching */
	if (*start && state.start == -1)
		_dat_walk_gfn_range(0, *start, asce, &ops, DAT_WALK_IGN_HOLES, &state);

	if (state.start == -1) {
		*count = 0;
	} else {
		*count = state.end - state.start;
		*start = state.start;
	}

	return 0;
}

struct dat_set_cmma_state {
	unsigned long mask;
	const u8 *bits;
};

static long __dat_set_cmma_pte(union pte *ptep, gfn_t gfn, gfn_t next, struct dat_walk *walk)
{
	struct dat_set_cmma_state *state = walk->priv;
	union pgste pgste, tmp;

	tmp.val = (state->bits[gfn - walk->start] << 24) & state->mask;

	pgste = pgste_get_lock(ptep);
	pgste.usage = tmp.usage;
	pgste.nodat = tmp.nodat;
	pgste_set_unlock(ptep, pgste);

	return 0;
}

/**
 * dat_set_cmma_bits() - Set CMMA bits for a range of guest pages.
 * @mc: Cache used for allocations.
 * @asce: The ASCE of the guest.
 * @gfn: The guest frame of the fist page whose CMMA bits are to set.
 * @count: How many pages need to be processed.
 * @mask: Which PGSTE bits should be set.
 * @bits: Points to an array with the CMMA attributes.
 *
 * This function sets the CMMA attributes for the given pages. If the input
 * buffer has zero length, no action is taken, otherwise the attributes are
 * set and the mm->context.uses_cmm flag is set.
 *
 * Each byte in @bits contains new values for bits 32-39 of the PGSTE.
 * Currently, only the fields NT and US are applied.
 *
 * Return: %0 in case of success, a negative error value otherwise.
 */
int dat_set_cmma_bits(struct kvm_s390_mmu_cache *mc, union asce asce, gfn_t gfn,
		      unsigned long count, unsigned long mask, const uint8_t *bits)
{
	const struct dat_walk_ops ops = { .pte_entry = __dat_set_cmma_pte, };
	struct dat_set_cmma_state state = { .mask = mask, .bits = bits, };
	union crste *crstep;
	union pte *ptep;
	gfn_t cur;
	int rc;

	for (cur = ALIGN_DOWN(gfn, _PAGE_ENTRIES); cur < gfn + count; cur += _PAGE_ENTRIES) {
		rc = dat_entry_walk(mc, cur, asce, DAT_WALK_ALLOC, TABLE_TYPE_PAGE_TABLE,
				    &crstep, &ptep);
		if (rc)
			return rc;
	}
	return _dat_walk_gfn_range(gfn, gfn + count, asce, &ops, DAT_WALK_IGN_HOLES, &state);
}