/*-
 * Copyright (c) 2019 Taylor R. Campbell
 * 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, 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 "cradix.h"

#ifdef _KERNEL

#include <sys/cdefs.h>
#include <sys/types.h>
#include <sys/atomic.h>
#include <sys/bitops.h>
#include <sys/kmem.h>

#define	assert	KASSERT

#else  /* !_KERNEL */

#include <sys/bitops.h>

#include <assert.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <strings.h>

#define	atomic_load_consume(p)		(*(p))
#define	atomic_store_release(p,v)	(*(p) = (v))
#define	atomic_store_relaxed(p,v)	(*(p) = (v))

#ifdef NDEBUG
#define	__diagused	__unused
#else
#define	__diagused
#endif

#endif	/* _KERNEL */

#define	arraycount(A)	(sizeof(A)/sizeof(*(A)))
#define	ffsl		ffs32	/* XXX */

#define	CHUNK_SIZE	CRADIX_CHUNK_SIZE
#define	CHUNK_MASK	CRADIX_CHUNK_MASK
#define	MAX_FANOUT	CRADIX_MAX_FANOUT
#define	NCHUNKS		CRADIX_NCHUNKS

struct cradix_branch {
	cradix_bitmap_t	bitmap;
	uint8_t		i;	/* key index */
	uint8_t		n;	/* number of children allocated */
	uintptr_t	child[];
};

#ifdef _KERNEL

static void *
cradix_default_alloc(size_t size)
{
	return kmem_alloc(size, KM_SLEEP);
}

static void
cradix_default_free(void *ptr, size_t size)
{
	return kmem_free(ptr, size);
}

static void
cradix_default_recycle(void *ptr, size_t size)
{
	return kmem_free(ptr, size);
}

#else

static void *
cradix_default_alloc(size_t size)
{
	return malloc(size);
}

static void
cradix_default_free(void *ptr, size_t size __unused)
{
	return free(ptr);
}

static void
cradix_default_recycle(void *ptr, size_t size __unused)
{
	return free(ptr);
}

#endif

static struct cradix_branch *
cradix_branch_alloc(struct cradix *C, unsigned n)
{
	size_t size = offsetof(struct cradix_branch, child[n]);
	struct cradix_branch *B;

	if (C->ops->alloc)
		B = (*C->ops->alloc)(C, size);
	else
		B = cradix_default_alloc(size);
	if (B == NULL)
		return NULL;
	B->n = n;
	return B;
}

static void
cradix_branch_free(struct cradix *C, struct cradix_branch *B)
{
	size_t size = offsetof(struct cradix_branch, child[B->n]);

	if (C->ops->free)
		(*C->ops->free)(C, B, size);
	else
		cradix_default_free(B, size);
}

static void
cradix_branch_recycle(struct cradix *C, struct cradix_branch *B)
{
	size_t size = offsetof(struct cradix_branch, child[B->n]);

	if (C->ops->recycle)
		(*C->ops->recycle)(C, B, size);
	else
		cradix_default_recycle(B, size);
}

void
cradix_init(struct cradix *C, const struct cradix_ops *ops)
{

	C->root = 0;
	C->ops = ops;
}

void
cradix_destroy(struct cradix *C)
{

	assert(C->root == 0);
	C->root = 2;		/* poison */
	C->ops = NULL;
}

static uint64_t
cradix_key(const struct cradix *C, const struct cradix_node *N)
{
	return (*C->ops->key)(C, N);
}

static unsigned
key_chunk(uint64_t key, unsigned i)
{
	/* Big-endian, so nearby addresses have nearby leaves.  */
	return (key >> ((NCHUNKS - i - 1)*CHUNK_SIZE)) & CHUNK_MASK;
}

static cradix_bitmap_t
map_chunk(cradix_bitmap_t bitmap, unsigned chunk)
{
	return bitmap | ((cradix_bitmap_t)1 << chunk);
}

static cradix_bitmap_t
unmap_chunk(cradix_bitmap_t bitmap, unsigned chunk)
{
	return bitmap & ~((cradix_bitmap_t)1 << chunk);
}

static bool
chunk_mapped(cradix_bitmap_t bitmap, unsigned chunk)
{
	return bitmap & ((cradix_bitmap_t)1 << chunk);
}

static unsigned
chunkpos(cradix_bitmap_t bitmap, unsigned chunk)
{
	return popcountl(bitmap & ~(~(cradix_bitmap_t)0 << chunk));
}

/* Tagged child pointers */

static inline bool
isbranch(uintptr_t child)
{

	return (child & 1) == 1;
}

static inline bool __diagused
isleaf(uintptr_t child)
{

	return (child & 1) == 0;
}

static inline uintptr_t
branch2child(struct cradix_branch *B)
{

	assert((uintptr_t)B != 0);
	assert(((uintptr_t)B & 1) == 0);

	return (uintptr_t)B | 1;
}

static inline uintptr_t
leaf2child(struct cradix_node *N)
{

	assert((uintptr_t)N != 0);
	assert(((uintptr_t)N & 1) == 0);

	return (uintptr_t)N;
}

static inline struct cradix_branch *
child2branch(uintptr_t child)
{
	struct cradix_branch *B;

	assert((child & 1) == 1);
	B = (struct cradix_branch *)(child - 1);
	assert(B != NULL);

	return B;
}

static inline struct cradix_node *
child2leaf(uintptr_t child)
{
	struct cradix_node *N;

	assert((child & 1) == 0);
	N = (struct cradix_node *)child;
	assert(N != NULL);

	return N;
}

struct cradix_node *
cradix_lookup(const struct cradix *C, uint64_t key)
{
	uintptr_t child;
	const struct cradix_branch *B;
	cradix_bitmap_t bitmap;
	uint32_t chunk, pos;
	struct cradix_node *N0;
	uint64_t key0;

	/* Load the root.  If there's nothing there, fail.  */
	if ((child = atomic_load_consume(&C->root)) == 0)
		return NULL;

	/* Find the best candidate match.  */
	for (;
	     isbranch(child);
	     child = atomic_load_consume(&B->child[pos])) {
		B = child2branch(child);
		chunk = key_chunk(key, B->i);
		bitmap = atomic_load_consume(&B->bitmap);
		if (!chunk_mapped(bitmap, chunk))
			return NULL;
		pos = chunkpos(bitmap, chunk);
	}

	/* Ignore tombstones.  */
	if (child == 0)
		return NULL;

	/* Verify whether it matches.  */
	assert(isleaf(child));
	N0 = child2leaf(child);
	key0 = cradix_key(C, N0);
	if (key0 != key)
		return NULL;

	return N0;
}

bool
cradix_insert_begin(struct cradix *C, struct cradix_insert *I)
{
	unsigned i;

	/* Preallocate one of each branch size we might need.  */
	for (i = 0; i < arraycount(I->alloc); i++) {
		unsigned n = 1u << (i + 1);
		if ((I->alloc[i] = cradix_branch_alloc(C, n)) == NULL)
			goto fail;
	}

	/* Initialize I->recycle to NULL.  */
	I->recycle = NULL;
	return true;

fail:	while (i --> 0) {
		cradix_branch_free(C, I->alloc[i]);
		I->alloc[i] = NULL; /* paranoia */
	}
	return false;
}

void
cradix_insert_end(struct cradix *C, struct cradix_insert *I)
{
	unsigned i;

	/* Free all the preallocated branches we didn't use.  */
	for (i = 0; i < arraycount(I->alloc); i++) {
		if (I->alloc[i])
			cradix_branch_free(C, I->alloc[i]);
		I->alloc[i] = NULL; /* paranoia */
	}

	/* If we have a node to recycle, schedule it for recycling.  */
	if (I->recycle) {
		cradix_branch_recycle(C, I->recycle);
		I->recycle = NULL; /* paranoia */
	}
}

static struct cradix_branch *
new_branch(struct cradix *C, struct cradix_insert *I)
{
	struct cradix_branch *B0;

	/* Allocate a 2-child branch.  */
	if (I) {
		B0 = I->alloc[0];
		assert(B0 != NULL);
		assert(B0->n >= 2);
		I->alloc[0] = NULL;
	} else {
		if ((B0 = cradix_branch_alloc(C, 2)) == NULL)
			return NULL;
	}

	return B0;
}

static struct cradix_branch *
grow_branch(struct cradix *C, struct cradix_insert *I,
    struct cradix_branch *B, unsigned new_chunk, uintptr_t new_child)
{
	struct cradix_branch *B0;
	cradix_bitmap_t bitmap, bitmap0;
	unsigned n0, chunk;

	/* Count the nondeleted children and assemble a bitmap of them.  */
	n0 = 0;
	bitmap0 = 0;
	bitmap = B->bitmap;
	while (bitmap) {
		chunk = ffsl(bitmap) - 1;
		assert(chunk != new_chunk);
		bitmap = unmap_chunk(bitmap, chunk);
		if (B->child[chunkpos(B->bitmap, chunk)] == 0)
			continue;
		n0++;
		bitmap0 = map_chunk(bitmap0, chunk);
	}
	assert(1 < n0);
	assert(n0 < CRADIX_MAX_FANOUT);
	assert(n0 <= B->n);

	/* Count the new child and enter it into the bitmap.  */
	n0++;
	bitmap0 = map_chunk(bitmap0, new_chunk);

	/* Allocate a branch -- possibly smaller than before!  */
	if (I) {
		unsigned i = ilog2(n0 - 1);
		B0 = I->alloc[i];
		assert(B0 != NULL);
		assert(B0->n >= n0);
		I->alloc[i] = NULL;
	} else {
		if ((B0 = cradix_branch_alloc(C, n0)) == NULL)
			return NULL;
	}
	B0->bitmap = bitmap0;
	B0->i = B->i;

	/* Copy over the existing nondeleted chunks.  */
	bitmap = B->bitmap;
	while (bitmap) {
		chunk = ffsl(bitmap) - 1;
		assert(chunk != new_chunk);
		bitmap = unmap_chunk(bitmap, chunk);
		if (B->child[chunkpos(B->bitmap, chunk)] == 0)
			continue;
		B0->child[chunkpos(B0->bitmap, chunk)] =
		    B->child[chunkpos(B->bitmap, chunk)];
	}

	/* Copy over the new chunk.  */
	B0->child[chunkpos(B0->bitmap, new_chunk)] = new_child;

	/* Recycle the old branch.  */
	if (I)
		I->recycle = B;
	else
		cradix_branch_recycle(C, B);

	return B0;
}

/*
 * Pick a leaf -- any leaf -- in the subtree rooted at B.
 */

static uintptr_t
pick_leaf(struct cradix_branch *B)
{
	cradix_bitmap_t bitmap;
	unsigned chunk;
	uintptr_t child, candidate = 0;

	/*
	 * Search breadth-first to try everything we can in the cache
	 * first.  No need for a queue because if we descend far enough
	 * we'll eventually hit a leaf.
	 *
	 * Conceivably we could store some metadata to reduce the time
	 * for this search, e.g. the maximum depth of a subtree, but
	 * it's not clear that maintaining the metadata is worth it.
	 */
	for (;;) {
		bitmap = B->bitmap;
		while (bitmap) {
			chunk = ffsl(bitmap) - 1;
			bitmap = unmap_chunk(bitmap, chunk);
			child = B->child[chunkpos(B->bitmap, chunk)];
			if (child == 0)
				continue;
			if (isleaf(child))
				return child;
			candidate = child;
		}
		assert(candidate);
		assert(isbranch(candidate));
		B = child2branch(candidate);
	}
}

struct cradix_node *
cradix_insert(struct cradix *C, struct cradix_node *N, struct cradix_insert *I)
{
	uint64_t key0, key = cradix_key(C, N);
	uintptr_t child, *childp;
	struct cradix_branch *B, *B0;
	unsigned chunk0, chunk, pos, i;
	struct cradix_node *N0;

	/* If the tree is empty, insert it at the root.  */
	if (C->root == 0) {
		atomic_store_release(&C->root, leaf2child(N));
		return N;
	}

	/* Find the first candidate match.  */
	for (child = C->root; isbranch(child); child = B->child[pos]) {
		B = child2branch(child);
		chunk = key_chunk(key, B->i);
		if (!chunk_mapped(B->bitmap, chunk)) {
			child = 0;
			break;
		}
		pos = chunkpos(B->bitmap, chunk);
	}

	/*
	 * If we hit a dead end, find any other node in the subtree as
	 * a candidate -- we know all nodes in this subtree agree on
	 * the all chunks up to the one we stopped at, so we just need
	 * any one of them to see whether we disagree earlier.
	 */
	if (child == 0)
		child = pick_leaf(B);
	assert(child != 0);

	/* Get the candidate node and key.  */
	assert(isleaf(child));
	N0 = child2leaf(child);
	key0 = cradix_key(C, N0);

	/* Find the index of the first mismatching chunk.  */
	for (i = 0; i < NCHUNKS; i++) {
		if (key_chunk(key0, i) != key_chunk(key, i))
			goto branch;
	}
	return N0;		/* collision */

branch:	/* Find where to insert a new branch and insert it.  */
	for (childp = &C->root; isbranch(*childp); childp = &B->child[pos]) {
		B = child2branch(*childp);
		if (i < B->i) {
			child = branch2child(B);
			break;
		}
		chunk = key_chunk(key, B->i);
		if (i == B->i)
			goto leaf;
		assert(chunk_mapped(B->bitmap, chunk));
		pos = chunkpos(B->bitmap, chunk);
	}

	/* Can't be a tombstone -- otherwise we would have stopped here.  */
	assert(*childp != 0);

	/* Splice a new branch at *childp.  */
	chunk = key_chunk(key, i);
	chunk0 = key_chunk(key0, i);
	if ((B0 = new_branch(C, I)) == NULL)
		return NULL;
	B0->i = i;
	B0->bitmap = 0;
	B0->bitmap = map_chunk(B0->bitmap, chunk);
	B0->bitmap = map_chunk(B0->bitmap, chunk0);
	B0->child[chunkpos(B0->bitmap, chunk)] = leaf2child(N);
	B0->child[chunkpos(B0->bitmap, chunk0)] = *childp;
	atomic_store_release(childp, branch2child(B0));
	return N;

leaf:	/*
	 * Grow a new leaf under this branch.  If the branch already
	 * has space for it filled with a tombstone, use that space;
	 * otherwise grow a larger branch.
	 */
	if (chunk_mapped(B->bitmap, chunk)) {
		pos = chunkpos(B->bitmap, chunk);
		assert(B->child[pos] == 0);
		atomic_store_release(&B->child[pos], leaf2child(N));
	} else {
		if ((B0 = grow_branch(C, I, B, chunk, leaf2child(N))) == NULL)
			return NULL;
		atomic_store_release(childp, branch2child(B0));
	}
	return N;
}

void
cradix_delete_begin(struct cradix *C __unused, struct cradix_delete *D)
{

	D->recycle = NULL;
}

void
cradix_delete_end(struct cradix *C, struct cradix_delete *D)
{

	if (D->recycle) {
		cradix_branch_recycle(C, D->recycle);
		D->recycle = NULL; /* paranoia */
	}
}

void
cradix_delete(struct cradix *C, struct cradix_node *N, struct cradix_delete *D)
{
	uint64_t key;
	struct cradix_node *N0;

	key = cradix_key(C, N);
	N0 = cradix_delete_key(C, key, D);
	assert(N0 == N);
}

static void
prune_branch(struct cradix *C, struct cradix_delete *D,
    struct cradix_branch *B, unsigned del_chunk, uintptr_t *nodep)
{
	cradix_bitmap_t bitmap, bitmap0;
	unsigned n0, chunk;

	assert(*nodep == branch2child(B));
	assert(chunk_mapped(B->bitmap, del_chunk));
	assert(B->child[chunkpos(B->bitmap, del_chunk)] != 0);

	/* Count the nondeleted children.  */
	n0 = 0;
	bitmap = B->bitmap;
	bitmap0 = 0;
	while (bitmap) {
		chunk = ffsl(bitmap) - 1;
		bitmap = unmap_chunk(bitmap, chunk);
		if (B->child[chunkpos(B->bitmap, chunk)] == 0)
			continue;
		bitmap0 = map_chunk(bitmap0, chunk);
		n0++;
	}
	assert(1 < n0);
	assert(n0 <= CRADIX_MAX_FANOUT);

	/* Check whether what's left is the last one.  */
	if (n0 == 2) {
		/* It's the last one -- splice and recycle the branch.  */
		bitmap0 = unmap_chunk(bitmap0, del_chunk);
		chunk = ffsl(bitmap0) - 1;
		atomic_store_relaxed(nodep,
		    B->child[chunkpos(B->bitmap, chunk)]);
		if (D)
			D->recycle = B;
		else
			cradix_branch_recycle(C, B);
	} else {
		/* There are siblings remaining.  Just leave a tombstone.  */
		atomic_store_relaxed(&B->child[chunkpos(B->bitmap, del_chunk)],
		    0);
	}
}

struct cradix_node *
cradix_delete_key(struct cradix *C, uint64_t key, struct cradix_delete *D)
{
	struct cradix_branch *B = NULL /*XXXGCC*/;
	struct cradix_node *N0;
	uintptr_t *parentp, *childp;
	uint64_t key0;
	unsigned chunk = -1 /*XXXGCC*/, pos = -1 /*XXXGCC*/;

	/* If there's nothing at the root, fail.  */
	if (C->root == 0)
		return NULL;

	/* Find the position of the best candidate match.  */
	for (parentp = NULL, childp = &C->root;
	     isbranch(*childp);
	     parentp = childp, childp = &B->child[pos]) {
		B = child2branch(*childp);
		chunk = key_chunk(key, B->i);
		if (!chunk_mapped(B->bitmap, chunk))
			return NULL;
		pos = chunkpos(B->bitmap, chunk);
	}

	/* If it's a tombstone, nothing to do.  */
	if (*childp == 0)
		return NULL;

	/* Verify whether it matches.  */
	assert(isleaf(*childp));
	N0 = child2leaf(*childp);
	key0 = cradix_key(C, N0);
	if (key0 != key)
		return NULL;

	/* Were we at the root or a branch?  */
	if (parentp == NULL) {
		/* Root -- just replace it by empty.  */
		assert(childp == &C->root);
		assert(C->root == leaf2child(N0));
		atomic_store_relaxed(&C->root, 0);
	} else {
		/* Branch -- prune it.  */
		assert(B != NULL);
		assert(B->child[pos] == leaf2child(N0));
		prune_branch(C, D, B, chunk, parentp);
	}

	/* Return the node we deleted.  */
	return N0;
}

#include <stdio.h>

static void
cradix_dump_child(const struct cradix *C, uintptr_t child, unsigned indent)
{

	if (child == 0) {
		printf("tombstone\n");
	} else if (isleaf(child)) {
		struct cradix_node *N = child2leaf(child);
		uint64_t key = cradix_key(C, N);

		printf("leaf @ %p: %016"PRIx64"\n", N, key);
	} else {
		struct cradix_branch *B = child2branch(child);
		cradix_bitmap_t bitmap;
		unsigned chunk, pos;
		uintptr_t child;

		printf("branch @ %p on chunk %u\n", B, B->i);
		indent++;

		bitmap = B->bitmap;
		while (bitmap) {
			chunk = ffsl(bitmap) - 1;
			bitmap = unmap_chunk(bitmap, chunk);
			pos = chunkpos(B->bitmap, chunk);
			child = atomic_load_consume(&B->child[pos]);

			printf("%*s[%02x] ", indent*2, "", chunk);
			cradix_dump_child(C, child, indent);
		}
	}
}

void
cradix_dump(const struct cradix *C)
{
	uintptr_t root;

	root = atomic_load_consume(&C->root);
	if (root == 0)
		printf("empty\n");
	else
		cradix_dump_child(C, root, 0);
}