/*-
 * Copyright (c) 2010 Isilon Systems, Inc.
 * Copyright (c) 2010 iX Systems, Inc.
 * Copyright (c) 2010 Panasas, Inc.
 * Copyright (c) 2013-2017 Mellanox Technologies, Ltd.
 * All rights reserved.
 *
 * 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 unmodified, 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 ``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 BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/kernel.h>
#include <sys/sysctl.h>
#include <sys/lock.h>
#include <sys/mutex.h>
#include <sys/stdarg.h>

#include <linux/bitmap.h>
#include <linux/kobject.h>
#include <linux/slab.h>
#include <linux/idr.h>
#include <linux/err.h>

#define	MAX_IDR_LEVEL	((MAX_IDR_SHIFT + IDR_BITS - 1) / IDR_BITS)
#define	MAX_IDR_FREE	(MAX_IDR_LEVEL * 2)

struct linux_idr_cache {
	spinlock_t lock;
	struct idr_layer *head;
	unsigned count;
};

DPCPU_DEFINE_STATIC(struct linux_idr_cache, linux_idr_cache);

/*
 * IDR Implementation.
 *
 * This is quick and dirty and not as re-entrant as the linux version
 * however it should be fairly fast.  It is basically a radix tree with
 * a builtin bitmap for allocation.
 */
static MALLOC_DEFINE(M_IDR, "idr", "Linux IDR compat");

static struct idr_layer *
idr_preload_dequeue_locked(struct linux_idr_cache *lic)
{
	struct idr_layer *retval;

	/* check if wrong thread is trying to dequeue */
	if (mtx_owned(&lic->lock) == 0)
		return (NULL);

	retval = lic->head;
	if (likely(retval != NULL)) {
		lic->head = retval->ary[0];
		lic->count--;
		retval->ary[0] = NULL;
	}
	return (retval);
}

static void
idr_preload_init(void *arg)
{
	int cpu;

	CPU_FOREACH(cpu) {
		struct linux_idr_cache *lic =
		    DPCPU_ID_PTR(cpu, linux_idr_cache);

		spin_lock_init(&lic->lock);
	}
}
SYSINIT(idr_preload_init, SI_SUB_CPU, SI_ORDER_ANY, idr_preload_init, NULL);

static void
idr_preload_uninit(void *arg)
{
	int cpu;

	CPU_FOREACH(cpu) {
		struct idr_layer *cacheval;
		struct linux_idr_cache *lic =
		    DPCPU_ID_PTR(cpu, linux_idr_cache);

		while (1) {
			spin_lock(&lic->lock);
			cacheval = idr_preload_dequeue_locked(lic);
			spin_unlock(&lic->lock);

			if (cacheval == NULL)
				break;
			free(cacheval, M_IDR);
		}
		spin_lock_destroy(&lic->lock);
	}
}
SYSUNINIT(idr_preload_uninit, SI_SUB_LOCK, SI_ORDER_FIRST, idr_preload_uninit, NULL);

void
idr_preload(gfp_t gfp_mask)
{
	struct linux_idr_cache *lic;
	struct idr_layer *cacheval;

	sched_pin();

	lic = &DPCPU_GET(linux_idr_cache);

	/* fill up cache */
	spin_lock(&lic->lock);
	while (lic->count < MAX_IDR_FREE) {
		spin_unlock(&lic->lock);
		cacheval = malloc(sizeof(*cacheval), M_IDR, M_ZERO | gfp_mask);
		spin_lock(&lic->lock);
		if (cacheval == NULL)
			break;
		cacheval->ary[0] = lic->head;
		lic->head = cacheval;
		lic->count++;
	}
}

void
idr_preload_end(void)
{
	struct linux_idr_cache *lic;

	lic = &DPCPU_GET(linux_idr_cache);
	spin_unlock(&lic->lock);
	sched_unpin();
}

static inline int
idr_max(struct idr *idr)
{
	return (1 << (idr->layers * IDR_BITS)) - 1;
}

static inline int
idr_pos(int id, int layer)
{
	return (id >> (IDR_BITS * layer)) & IDR_MASK;
}

void
idr_init(struct idr *idr)
{
	bzero(idr, sizeof(*idr));
	mtx_init(&idr->lock, "idr", NULL, MTX_DEF);
}

/* Only frees cached pages. */
void
idr_destroy(struct idr *idr)
{
	struct idr_layer *il, *iln;

	/*
	 * This idr can be reused, and this function might be called multiple times
	 * without a idr_init(). Check if this is the case.  If we do not do this
	 * then the mutex will panic while asserting that it is valid.
	 */
	if (mtx_initialized(&idr->lock) == 0)
		return;

	idr_remove_all(idr);
	mtx_lock(&idr->lock);
	for (il = idr->free; il != NULL; il = iln) {
		iln = il->ary[0];
		free(il, M_IDR);
	}
	mtx_unlock(&idr->lock);
	mtx_destroy(&idr->lock);
}

static void
idr_remove_layer(struct idr_layer *il, int layer)
{
	int i;

	if (il == NULL)
		return;
	if (layer == 0) {
		free(il, M_IDR);
		return;
	}
	for (i = 0; i < IDR_SIZE; i++)
		if (il->ary[i])
			idr_remove_layer(il->ary[i], layer - 1);
}

void
idr_remove_all(struct idr *idr)
{

	mtx_lock(&idr->lock);
	idr_remove_layer(idr->top, idr->layers - 1);
	idr->top = NULL;
	idr->layers = 0;
	mtx_unlock(&idr->lock);
}

static void *
idr_remove_locked(struct idr *idr, int id)
{
	struct idr_layer *il;
	void *res;
	int layer;
	int idx;

	id &= MAX_ID_MASK;
	il = idr->top;
	layer = idr->layers - 1;
	if (il == NULL || id > idr_max(idr))
		return (NULL);
	/*
	 * Walk down the tree to this item setting bitmaps along the way
	 * as we know at least one item will be free along this path.
	 */
	while (layer && il) {
		idx = idr_pos(id, layer);
		il->bitmap |= 1 << idx;
		il = il->ary[idx];
		layer--;
	}
	idx = id & IDR_MASK;
	/*
	 * At this point we've set free space bitmaps up the whole tree.
	 * We could make this non-fatal and unwind but linux dumps a stack
	 * and a warning so I don't think it's necessary.
	 */
	if (il == NULL || (il->bitmap & (1 << idx)) != 0)
		panic("idr_remove: Item %d not allocated (%p, %p)\n",
		    id, idr, il);
	res = il->ary[idx];
	il->ary[idx] = NULL;
	il->bitmap |= 1 << idx;

	return (res);
}

void *
idr_remove(struct idr *idr, int id)
{
	void *res;

	mtx_lock(&idr->lock);
	res = idr_remove_locked(idr, id);
	mtx_unlock(&idr->lock);

	return (res);
}

static inline struct idr_layer *
idr_find_layer_locked(struct idr *idr, int id)
{
	struct idr_layer *il;
	int layer;

	id &= MAX_ID_MASK;
	il = idr->top;
	layer = idr->layers - 1;
	if (il == NULL || id > idr_max(idr))
		return (NULL);
	while (layer && il) {
		il = il->ary[idr_pos(id, layer)];
		layer--;
	}
	return (il);
}

void *
idr_replace(struct idr *idr, void *ptr, int id)
{
	struct idr_layer *il;
	void *res;
	int idx;

	mtx_lock(&idr->lock);
	il = idr_find_layer_locked(idr, id);
	idx = id & IDR_MASK;

	/* Replace still returns an error if the item was not allocated. */
	if (il == NULL || (il->bitmap & (1 << idx))) {
		res = ERR_PTR(-ENOENT);
	} else {
		res = il->ary[idx];
		il->ary[idx] = ptr;
	}
	mtx_unlock(&idr->lock);
	return (res);
}

static inline void *
idr_find_locked(struct idr *idr, int id)
{
	struct idr_layer *il;
	void *res;

	mtx_assert(&idr->lock, MA_OWNED);
	il = idr_find_layer_locked(idr, id);
	if (il != NULL)
		res = il->ary[id & IDR_MASK];
	else
		res = NULL;
	return (res);
}

void *
idr_find(struct idr *idr, int id)
{
	void *res;

	mtx_lock(&idr->lock);
	res = idr_find_locked(idr, id);
	mtx_unlock(&idr->lock);
	return (res);
}

void *
idr_get_next(struct idr *idr, int *nextidp)
{
	void *res = NULL;
	int id = *nextidp;

	mtx_lock(&idr->lock);
	for (; id <= idr_max(idr); id++) {
		res = idr_find_locked(idr, id);
		if (res == NULL)
			continue;
		*nextidp = id;
		break;
	}
	mtx_unlock(&idr->lock);
	return (res);
}

int
idr_pre_get(struct idr *idr, gfp_t gfp_mask)
{
	struct idr_layer *il, *iln;
	struct idr_layer *head;
	int need;

	mtx_lock(&idr->lock);
	for (;;) {
		need = idr->layers + 1;
		for (il = idr->free; il != NULL; il = il->ary[0])
			need--;
		mtx_unlock(&idr->lock);
		if (need <= 0)
			break;
		for (head = NULL; need; need--) {
			iln = malloc(sizeof(*il), M_IDR, M_ZERO | gfp_mask);
			if (iln == NULL)
				break;
			bitmap_fill(&iln->bitmap, IDR_SIZE);
			if (head != NULL) {
				il->ary[0] = iln;
				il = iln;
			} else
				head = il = iln;
		}
		if (head == NULL)
			return (0);
		mtx_lock(&idr->lock);
		il->ary[0] = idr->free;
		idr->free = head;
	}
	return (1);
}

static struct idr_layer *
idr_free_list_get(struct idr *idp)
{
	struct idr_layer *il;

	if ((il = idp->free) != NULL) {
		idp->free = il->ary[0];
		il->ary[0] = NULL;
	}
	return (il);
}

static inline struct idr_layer *
idr_get(struct idr *idp)
{
	struct idr_layer *il;

	if ((il = idr_free_list_get(idp)) != NULL) {
		MPASS(il->bitmap != 0);
	} else if ((il = malloc(sizeof(*il), M_IDR, M_ZERO | M_NOWAIT)) != NULL) {
		bitmap_fill(&il->bitmap, IDR_SIZE);
	} else if ((il = idr_preload_dequeue_locked(&DPCPU_GET(linux_idr_cache))) != NULL) {
		bitmap_fill(&il->bitmap, IDR_SIZE);
	} else {
		return (NULL);
	}
	return (il);
}

/*
 * Could be implemented as get_new_above(idr, ptr, 0, idp) but written
 * first for simplicity sake.
 */
static int
idr_get_new_locked(struct idr *idr, void *ptr, int *idp)
{
	struct idr_layer *stack[MAX_LEVEL];
	struct idr_layer *il;
	int error;
	int layer;
	int idx;
	int id;

	mtx_assert(&idr->lock, MA_OWNED);

	error = -EAGAIN;
	/*
	 * Expand the tree until there is free space.
	 */
	if (idr->top == NULL || idr->top->bitmap == 0) {
		if (idr->layers == MAX_LEVEL + 1) {
			error = -ENOSPC;
			goto out;
		}
		il = idr_get(idr);
		if (il == NULL)
			goto out;
		il->ary[0] = idr->top;
		if (idr->top)
			il->bitmap &= ~1;
		idr->top = il;
		idr->layers++;
	}
	il = idr->top;
	id = 0;
	/*
	 * Walk the tree following free bitmaps, record our path.
	 */
	for (layer = idr->layers - 1;; layer--) {
		stack[layer] = il;
		idx = ffsl(il->bitmap);
		if (idx == 0)
			panic("idr_get_new: Invalid leaf state (%p, %p)\n",
			    idr, il);
		idx--;
		id |= idx << (layer * IDR_BITS);
		if (layer == 0)
			break;
		if (il->ary[idx] == NULL) {
			il->ary[idx] = idr_get(idr);
			if (il->ary[idx] == NULL)
				goto out;
		}
		il = il->ary[idx];
	}
	/*
	 * Allocate the leaf to the consumer.
	 */
	il->bitmap &= ~(1 << idx);
	il->ary[idx] = ptr;
	*idp = id;
	/*
	 * Clear bitmaps potentially up to the root.
	 */
	while (il->bitmap == 0 && ++layer < idr->layers) {
		il = stack[layer];
		il->bitmap &= ~(1 << idr_pos(id, layer));
	}
	error = 0;
out:
#ifdef INVARIANTS
	if (error == 0 && idr_find_locked(idr, id) != ptr) {
		panic("idr_get_new: Failed for idr %p, id %d, ptr %p\n",
		    idr, id, ptr);
	}
#endif
	return (error);
}

int
idr_get_new(struct idr *idr, void *ptr, int *idp)
{
	int retval;

	mtx_lock(&idr->lock);
	retval = idr_get_new_locked(idr, ptr, idp);
	mtx_unlock(&idr->lock);
	return (retval);
}

static int
idr_get_new_above_locked(struct idr *idr, void *ptr, int starting_id, int *idp)
{
	struct idr_layer *stack[MAX_LEVEL];
	struct idr_layer *il;
	int error;
	int layer;
	int idx, sidx;
	int id;

	mtx_assert(&idr->lock, MA_OWNED);

	error = -EAGAIN;
	/*
	 * Compute the layers required to support starting_id and the mask
	 * at the top layer.
	 */
restart:
	idx = starting_id;
	layer = 0;
	while (idx & ~IDR_MASK) {
		layer++;
		idx >>= IDR_BITS;
	}
	if (layer == MAX_LEVEL + 1) {
		error = -ENOSPC;
		goto out;
	}
	/*
	 * Expand the tree until there is free space at or beyond starting_id.
	 */
	while (idr->layers <= layer ||
	    idr->top->bitmap < (1 << idr_pos(starting_id, idr->layers - 1))) {
		if (idr->layers == MAX_LEVEL + 1) {
			error = -ENOSPC;
			goto out;
		}
		il = idr_get(idr);
		if (il == NULL)
			goto out;
		il->ary[0] = idr->top;
		if (idr->top && idr->top->bitmap == 0)
			il->bitmap &= ~1;
		idr->top = il;
		idr->layers++;
	}
	il = idr->top;
	id = 0;
	/*
	 * Walk the tree following free bitmaps, record our path.
	 */
	for (layer = idr->layers - 1;; layer--) {
		stack[layer] = il;
		sidx = idr_pos(starting_id, layer);
		/* Returns index numbered from 0 or size if none exists. */
		idx = find_next_bit(&il->bitmap, IDR_SIZE, sidx);
		if (idx == IDR_SIZE && sidx == 0)
			panic("idr_get_new: Invalid leaf state (%p, %p)\n",
			    idr, il);
		/*
		 * We may have walked a path where there was a free bit but
		 * it was lower than what we wanted.  Restart the search with
		 * a larger starting id.  id contains the progress we made so
		 * far.  Search the leaf one above this level.  This may
		 * restart as many as MAX_LEVEL times but that is expected
		 * to be rare.
		 */
		if (idx == IDR_SIZE) {
			starting_id = id + (1 << ((layer + 1) * IDR_BITS));
			goto restart;
		}
		if (idx > sidx)
			starting_id = 0;	/* Search the whole subtree. */
		id |= idx << (layer * IDR_BITS);
		if (layer == 0)
			break;
		if (il->ary[idx] == NULL) {
			il->ary[idx] = idr_get(idr);
			if (il->ary[idx] == NULL)
				goto out;
		}
		il = il->ary[idx];
	}
	/*
	 * Allocate the leaf to the consumer.
	 */
	il->bitmap &= ~(1 << idx);
	il->ary[idx] = ptr;
	*idp = id;
	/*
	 * Clear bitmaps potentially up to the root.
	 */
	while (il->bitmap == 0 && ++layer < idr->layers) {
		il = stack[layer];
		il->bitmap &= ~(1 << idr_pos(id, layer));
	}
	error = 0;
out:
#ifdef INVARIANTS
	if (error == 0 && idr_find_locked(idr, id) != ptr) {
		panic("idr_get_new_above: Failed for idr %p, id %d, ptr %p\n",
		    idr, id, ptr);
	}
#endif
	return (error);
}

int
idr_get_new_above(struct idr *idr, void *ptr, int starting_id, int *idp)
{
	int retval;

	mtx_lock(&idr->lock);
	retval = idr_get_new_above_locked(idr, ptr, starting_id, idp);
	mtx_unlock(&idr->lock);
	return (retval);
}

int
ida_get_new_above(struct ida *ida, int starting_id, int *p_id)
{
	return (idr_get_new_above(&ida->idr, NULL, starting_id, p_id));
}

static int
idr_alloc_locked(struct idr *idr, void *ptr, int start, int end)
{
	int max = end > 0 ? end - 1 : INT_MAX;
	int error;
	int id;

	mtx_assert(&idr->lock, MA_OWNED);

	if (unlikely(start < 0))
		return (-EINVAL);
	if (unlikely(max < start))
		return (-ENOSPC);

	if (start == 0)
		error = idr_get_new_locked(idr, ptr, &id);
	else
		error = idr_get_new_above_locked(idr, ptr, start, &id);

	if (unlikely(error < 0))
		return (error);
	if (unlikely(id > max)) {
		idr_remove_locked(idr, id);
		return (-ENOSPC);
	}
	return (id);
}

int
idr_alloc(struct idr *idr, void *ptr, int start, int end, gfp_t gfp_mask)
{
	int retval;

	mtx_lock(&idr->lock);
	retval = idr_alloc_locked(idr, ptr, start, end);
	mtx_unlock(&idr->lock);
	return (retval);
}

int
idr_alloc_cyclic(struct idr *idr, void *ptr, int start, int end, gfp_t gfp_mask)
{
	int retval;

	mtx_lock(&idr->lock);
	retval = idr_alloc_locked(idr, ptr, max(start, idr->next_cyclic_id), end);
	if (unlikely(retval == -ENOSPC))
		retval = idr_alloc_locked(idr, ptr, start, end);
	if (likely(retval >= 0))
		idr->next_cyclic_id = retval + 1;
	mtx_unlock(&idr->lock);
	return (retval);
}

static int
idr_for_each_layer(struct idr_layer *il, int offset, int layer,
    int (*f)(int id, void *p, void *data), void *data)
{
	int i, err;

	if (il == NULL)
		return (0);
	if (layer == 0) {
		for (i = 0; i < IDR_SIZE; i++) {
			if (il->ary[i] == NULL)
				continue;
			err = f(i + offset, il->ary[i],  data);
			if (err)
				return (err);
		}
		return (0);
	}
	for (i = 0; i < IDR_SIZE; i++) {
		if (il->ary[i] == NULL)
			continue;
		err = idr_for_each_layer(il->ary[i],
		    (i + offset) * IDR_SIZE, layer - 1, f, data);
		if (err)
			return (err);
	}
	return (0);
}

/* NOTE: It is not allowed to modify the IDR tree while it is being iterated */
int
idr_for_each(struct idr *idp, int (*f)(int id, void *p, void *data), void *data)
{
	return (idr_for_each_layer(idp->top, 0, idp->layers - 1, f, data));
}

static int
idr_has_entry(int id, void *p, void *data)
{

	return (1);
}

bool
idr_is_empty(struct idr *idp)
{

	return (idr_for_each(idp, idr_has_entry, NULL) == 0);
}

int
ida_pre_get(struct ida *ida, gfp_t flags)
{
	if (idr_pre_get(&ida->idr, flags) == 0)
		return (0);

	if (ida->free_bitmap == NULL) {
		ida->free_bitmap =
		    malloc(sizeof(struct ida_bitmap), M_IDR, flags);
	}
	return (ida->free_bitmap != NULL);
}

int
ida_simple_get(struct ida *ida, unsigned int start, unsigned int end,
    gfp_t flags)
{
	int ret, id;
	unsigned int max;

	MPASS((int)start >= 0);

	if ((int)end <= 0)
		max = INT_MAX;
	else {
		MPASS(end > start);
		max = end - 1;
	}
again:
	if (!ida_pre_get(ida, flags))
		return (-ENOMEM);

	if ((ret = ida_get_new_above(ida, start, &id)) == 0) {
		if (id > max) {
			ida_remove(ida, id);
			ret = -ENOSPC;
		} else {
			ret = id;
		}
	}
	if (__predict_false(ret == -EAGAIN))
		goto again;

	return (ret);
}

void
ida_simple_remove(struct ida *ida, unsigned int id)
{
	idr_remove(&ida->idr, id);
}

void
ida_remove(struct ida *ida, int id)
{
	idr_remove(&ida->idr, id);
}

void
ida_init(struct ida *ida)
{
	idr_init(&ida->idr);
}

void
ida_destroy(struct ida *ida)
{
	idr_destroy(&ida->idr);
	free(ida->free_bitmap, M_IDR);
	ida->free_bitmap = NULL;
}