/*
 * Copyright 2016 Red Hat Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * Authors: Jérôme Glisse <jglisse@redhat.com>
 */
/*
 * Generic page table structure with adjustable depth. For details refer to
 * include/linux/gpt.h
 */
#include <linux/mm.h>
#include <linux/gpt.h>
#include <linux/slab.h>
#include <linux/highmem.h>

#ifdef CONFIG_HIGHMEM64G
/* Some arch do not define MAX_PHYSMEM_BITS */
#ifndef MAX_PHYSMEM_BITS
#define MAX_PHYSMEM_BITS 36
#endif /* MAX_PHYSMEM_BITS */
#define GTE_SHIFT 3
#define GTE_BITS 64
/* GPT_GTD_SHIFT - Shift for one directory level index */
#else /* CONFIG_HIGHMEM64G */
/* Some arch do not define MAX_PHYSMEM_BITS */
#ifndef MAX_PHYSMEM_BITS
#define MAX_PHYSMEM_BITS BITS_PER_LONG
#endif /* MAX_PHYSMEM_BITS */
#if BITS_PER_LONG == 32
#define GTE_SHIFT 2
#define GTE_BITS 32
#else /* This must be 64 McFly ! */
#define GTE_SHIFT 3
#define GTE_BITS 64
#endif
#endif /* CONFIG_HIGHMEM64G */

#define GPT_DEFAULT_GFP (GFP_KERNEL | __GFP_ZERO | __GFP_HIGHMEM)


/*
 * GPT_PFN_BITS - Number of bits require to store any valid pfn value
 * GPT_PFN_MASK - Mask for pfn value
 *
 * This means that there is (BITS_PER_LONG - GPT_PFN_BITS) bits that can be use
 * in anyway by the end user of GPT struct.
 */
#define GPT_PFN_BITS (MAX_PHYSMEM_BITS - PAGE_SHIFT)
#define GPT_PFN_MASK ((1UL << GPT_PFN_BITS) - 1)

/*
 * GPT_GTD_SHIFT - Shift for one directory level index
 * GPT_GTE_PER_GTD - Number of page table entry in one directory level
 */
#define GTD_SHIFT (PAGE_SHIFT - GTE_SHIFT)
#define GTE_PER_GTD (1UL << GTD_SHIFT)
#define GTD_MASK (GTE_PER_GTD - 1)


struct gpt *gpt_alloc(unsigned long start,
		      unsigned long end,
		      uint8_t shift,
		      uint8_t valid_bit)
{
	unsigned long ngd, tmp;
	struct gpt *gpt;

	/* Sanity checks */
	start &= PAGE_MASK;
	end &= PAGE_MASK;
	BUG_ON(start >= end);
	BUG_ON(valid_bit >= shift);
	BUG_ON((GTE_BITS - shift) < GPT_PFN_BITS);

	gpt = kmalloc(sizeof(*gpt), GFP_KERNEL);
	if (!gpt)
		return NULL;

	gpt->start = start;
	gpt->end = end;
	gpt->nlevels = 0;
	gpt->shift = shift;
	gpt->valid_bit = valid_bit;
	spin_lock_init(&gpt->lock);
	atomic_set(&gpt->refcount, 0);

	ngd = (end - start) >> PAGE_SHIFT;
	tmp = (ngd) >> GTD_SHIFT;
	while (tmp) {
		ngd = tmp;
		tmp = tmp >> GTD_SHIFT;
		gpt->nlevels++;
	}
	BUG_ON(gpt->nlevels >= GPT_MAX_LEVEL);

	gpt->gdp = kzalloc(ngd * sizeof(gte_t), GFP_KERNEL);
	if (!gpt->gdp) {
		kfree(gpt);
		return NULL;
	}

	return gpt;
}
EXPORT_SYMBOL(gpt_alloc);

void gpt_free(struct gpt *gpt)
{
	struct gpt_walk walk;

	if (!gpt)
		return;

	gpt_walk_init(&walk, gpt);
	gpt_walk_prune(&walk, gpt->start, gpt->end);
	gpt_walk_fini(&walk);

	kfree(gpt->gdp);
	kfree(gpt);
}
EXPORT_SYMBOL(gpt_free);


static inline unsigned long gte_to_pfn(struct gpt_walk *walk, gte_t gte)
{
	return (unsigned long)(gte >> walk->gpt->shift) & GPT_PFN_MASK;
}

static inline struct page *gte_to_page(struct gpt_walk *walk, gte_t gte)
{
	if (!(gte & (1UL << walk->gpt->valid_bit)))
		return NULL;
	return pfn_to_page(gte_to_pfn(walk, gte));
}

static inline unsigned gpt_index(struct gpt_walk *walk,
				 unsigned long addr,
				 unsigned level)
{
	unsigned shift;

	shift = GTD_SHIFT * level + PAGE_SHIFT;
	return ((addr - walk->gpt->start) >> shift) & GTD_MASK;
}

static inline unsigned long gpt_level_start(struct gpt_walk *walk,
					    unsigned long addr,
					    unsigned level)
{
	unsigned long mask;
	unsigned shift;

	shift = GTD_SHIFT * (level + 1) + PAGE_SHIFT;
	mask = ~((1UL << shift) - 1);
	return ((addr - walk->gpt->start) & mask) + walk->gpt->start;
}

static inline unsigned long gpt_level_end(struct gpt_walk *walk,
					  unsigned long addr,
					  unsigned level)
{
	unsigned long mask;
	unsigned shift;

	shift = GTD_SHIFT * (level + 1) + PAGE_SHIFT;
	mask = (1UL << shift) - 1;
	return ((addr - walk->gpt->start) | mask) + walk->gpt->start + 1;
}


/* gpt_walk_init() - Init gpt walk structure
 * @walk: walk structure to initialize
 */
void gpt_walk_init(struct gpt_walk *walk, struct gpt *gpt)
{
	unsigned level;

	BUG_ON(!gpt);

	walk->gpt = gpt;
	walk->gtd[walk->gpt->nlevels] = NULL;
	walk->gte[walk->gpt->nlevels] = gpt->gdp;
	walk->start = gpt->start;
	walk->end = gpt->end;

	for (level = 0; level < walk->gpt->nlevels; level++) {
		walk->gtd[level] = NULL;
		walk->gte[level] = NULL;
	}
}
EXPORT_SYMBOL(gpt_walk_init);

/* gpt_walk_fini() - Finalize gpt walk structure
 * @walk: walk structure to finalize
 *
 * This unmap any mapped directory.
 */
void gpt_walk_fini(struct gpt_walk *walk)
{
	unsigned level;

	for (level = 0; level < walk->gpt->nlevels; ++level) {
		if (!walk->gtd[level])
			continue;
		kunmap(walk->gtd[level]);
		atomic_dec(gpt_walk_gtd_refcount(walk, level));
		walk->gtd[level] = NULL;
		walk->gte[level] = NULL;
	}
	walk->start = walk->gpt->start;
	walk->end = walk->gpt->end;
}
EXPORT_SYMBOL(gpt_walk_fini);

/* gpt_walk_gtep_from_addr() - Get entry pointer for a given address
 * @walk: walk structure use to walk generic page table
 * @addr: address of interest
 * Returns: NULL if address is not in range of if there is no directory
 *
 * This will return pointer to page table entry if a directory exist for the
 * given address.
 */
gte_t *gpt_walk_gtep_from_addr(struct gpt_walk *walk, unsigned long addr)
{
	unsigned l;

	if (addr < walk->gpt->start || addr >= walk->gpt->end)
		return NULL;

again:
	if (walk->gte[0] && addr >= walk->start && addr < walk->end)
		return &walk->gte[0][gpt_index(walk, addr, 0)];

	for (l = 0; l < walk->gpt->nlevels; l++) {
		if (!walk->gtd[l])
			continue;

		if (addr >= walk->start && addr < walk->end)
			break;

		kunmap(walk->gtd[l]);
		atomic_dec(gpt_walk_gtd_refcount(walk, l));
		walk->gtd[l] = NULL;
		walk->gte[l] = NULL;
		/* Compute start and end address of upper level */
		walk->start = gpt_level_start(walk, walk->start, l + 1);
		walk->end = gpt_level_end(walk, walk->start, l + 1);
	}

	for (; l; l--) {
		unsigned idx;
		atomic_t *refcount;

		idx = gpt_index(walk, addr, l);
		spin_lock(gpt_walk_gtd_lock_ptr(walk, l));
		walk->gtd[l - 1] = gte_to_page(walk, walk->gte[l][idx]);
		refcount = gpt_walk_gtd_refcount(walk, l - 1);
		if (refcount)
			atomic_inc(refcount);
		else
			walk->gtd[l - 1] = NULL;
		spin_unlock(gpt_walk_gtd_lock_ptr(walk, l));
		if (!walk->gtd[l- 1])
			return NULL;

		walk->gte[l - 1] = kmap(walk->gtd[l - 1]);

		/* Compute start and end address of lower level */
		walk->start = gpt_level_start(walk, addr, l - 1);
		walk->end = gpt_level_end(walk, addr, l - 1);
	}

	/* At this point all gtd levels are mapped */
	goto again;
}
EXPORT_SYMBOL(gpt_walk_gtep_from_addr);

/* gpt_populate() - Populate page table directory tree for given address
 * @walk: walk structure use to walk generic page table
 * @addr: address of interest
 * Returns: NULL if nothing to populate (locking error) directory entry pointer
 *
 * This will populate all directory levels that are missing for a given address
 * and it returns a pointer to lowest directory level.
 */
gte_t *gpt_walk_populate(struct gpt_walk *walk, unsigned long addr)
{
	unsigned idx;
	gte_t *gtep;
	int level;

	if ((gtep = gpt_walk_gtep_from_addr(walk, addr)))
		return gtep;

	for (level = walk->gpt->nlevels - 1; level >= 0; level--) {
		unsigned long pfn;

		if (walk->gtd[level])
			continue;

		walk->gtd[level] = alloc_page(GPT_DEFAULT_GFP);
		if (!walk->gtd[level])
			return NULL;
		if (!ptlock_init(walk->gtd[level])) {
			__free_page(walk->gtd[level]);
			walk->gtd[level] = NULL;
			return NULL;
		}
		pfn = page_to_pfn(walk->gtd[level]);

		/* Initialize new directory */
		atomic_set(&walk->gtd[level]->_mapcount, 1);

		/* Compute start and end address of current level */
		walk->start = gpt_level_start(walk, addr, level);
		walk->end = gpt_level_end(walk, addr, level);
		walk->gte[level] = kmap(walk->gtd[level]);

		/* Set directory entry in upper level */
		idx = gpt_index(walk, addr, level + 1);
		atomic_inc(gpt_walk_gtd_refcount(walk, level + 1));
		spin_lock(gpt_walk_gtd_lock_ptr(walk, level + 1));
		walk->gte[level + 1][idx] = pfn << walk->gpt->shift;
		walk->gte[level + 1][idx] |= 1UL << walk->gpt->valid_bit;
		spin_unlock(gpt_walk_gtd_lock_ptr(walk, level + 1));
	}

	idx = gpt_index(walk, addr, 0);
	return &walk->gte[0][idx];
}
EXPORT_SYMBOL(gpt_walk_populate);

/* gpt_page_reset() - Reset struct page fields use by gpt to sane value
 * @page: page that is being prune from directory tree and needs reset
 *
 * We use few fields inside the struct page to store informations for each of
 * directory in the tree. We need to reset some of those fields to sane value
 * so that core mm does not freaks out when we free a directory page.
 *
 * This should be call by the prune callback provided to gpt_prune() when it
 * decides to free a directory level. Default callback, gpt_prune_default(),
 * properly call this function.
 */
static inline void gpt_page_reset(struct page *page)
{
	atomic_set(&page->_mapcount, -1);
}

/* gpt_prune() - Prune page table directory tree for given address range
 * @walk: walk structure use to walk generic page table
 * @start: range start address
 * @end: range end address
 *
 * This will prune the directory tree for the given address range. Any empty
 * directory will be free.
 *
 * WARNING YOU ARE RESPONSIBLE FOR LOCKING IN RESPECT TO CONCURRENT PAGE TABLE
 * PRUNING OR POPULATING !
 */
void gpt_walk_prune(struct gpt_walk *walk,
		    unsigned long start,
		    unsigned long end)
{
	unsigned long addr;

	start &= PAGE_MASK;
	end &= PAGE_MASK;
	BUG_ON(start >= end);

	for (addr = start; addr < end;) {
		unsigned long next;
		unsigned l;

		next = min(end, gpt_level_end(walk, addr, 0));
		gpt_walk_gtep_from_addr(walk, addr);
		for (l = 0; l < walk->gpt->nlevels; l++) {
			unsigned idx;

			if (!walk->gtd[l]) {
				next = min(end, gpt_level_end(walk, addr, l));
				continue;
			}

			if (atomic_read(gpt_walk_gtd_refcount(walk, l)) != 1)
				break;

			/* First unmap end update walk structure */
			kunmap(walk->gtd[l]);
			walk->gtd[l] = NULL;
			walk->gte[l] = NULL;

			/* Pointer to directory entry in the upper level */
			idx = gpt_index(walk, walk->start, l + 1);
			spin_lock(gpt_walk_gtd_lock_ptr(walk, l + 1));
			walk->gte[l + 1][idx] = 0;
			spin_unlock(gpt_walk_gtd_lock_ptr(walk, l + 1));
			atomic_dec(gpt_walk_gtd_refcount(walk, l + 1));

			/* The next address is end address of current level */
			next = min(end, walk->end);

			/* Start and end address are now for the upper level */
			walk->start = gpt_level_start(walk, addr, l + 1);
			walk->end = gpt_level_end(walk, addr, l + 1);
		}
		addr = next;
	}
}
EXPORT_SYMBOL(gpt_walk_prune);

int gpt_walk_range(struct gpt_walk *walk,
		   unsigned long start,
		   unsigned long end,
		   gpt_walk_cb_t cb,
		   void *private)
{
	unsigned long addr;

	for (addr = start; addr < end;) {
		unsigned long next;
		spinlock_t *gtl;
		gte_t *gtep;
		int ret;

		gtep = gpt_walk_gtep_from_addr(walk, addr);
		if (!gtep) {
			unsigned l;

			for (l = 0; l < walk->gpt->nlevels; l++) {
				if (walk->gtd[l])
					break;
				addr = min(end, gpt_level_end(walk, addr, l));
			}
			continue;
		}

		next = min(end, gpt_level_end(walk, addr, 0));
		gtl = gpt_walk_gtd_lock_ptr(walk, 0);
		ret = cb(walk, addr, next, gtl, gtep, private);
		if (ret)
			return ret;
		addr = next;
	}

	return 0;
}
EXPORT_SYMBOL(gpt_walk_range);