/*-
 * Copyright (c) 2015--2017 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 <assert.h>
#include <stddef.h>
#include <stdint.h>

#include "critbit.h"

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

#if CRITBIT_DEBUG
#define	__dbgused
#else
#define	__dbgused	__unused
#endif

#if CRITBIT_DEBUG
# define	CRITBIT_BRANCH_MAGIC	UINT64_C(0x80462106d7c49800)
# define	CRITBIT_NODE_MAGIC	UINT64_C(0xae2e9ea4dac95b66)
# define	DBG(x)			x
#else
# define	DBG(x)			((void)0)
#endif

#define	MIN(x, y)		((x) < (y) ? (x) : (y))

/* Tagged child pointers */

static bool
isbranch(uintptr_t child)
{

	return (child & 1) == 1;
}

static bool __diagused
isleaf(uintptr_t child)
{

	return (child & 1) == 0;
}

static uintptr_t
branch2child(struct critbit_branch *B)
{

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

	return (uintptr_t)B | 1;
}

static uintptr_t
leaf2child(struct critbit_node *N)
{

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

	return (uintptr_t)N;
}

static struct critbit_branch *
child2branch(uintptr_t child)
{
	struct critbit_branch *B;

	assert((child & 1) == 1);
	B = (struct critbit_branch *)(child - 1);
	assert(B != NULL);
	DBG(assert(B->cbb_magic == CRITBIT_BRANCH_MAGIC));

	return B;
}

static struct critbit_node *
child2leaf(uintptr_t child)
{
	struct critbit_node *N;

	assert((child & 1) == 0);
	N = (struct critbit_node *)child;
	assert(N != NULL);
	DBG(assert(N->cbn_magic == CRITBIT_NODE_MAGIC));

	return N;
}

/* Critbit tree and node creation and destruction */

void
critbit__init(struct critbit *C, const struct critbit_ops *ops,
    int debug __diagused)
{

#if CRITBIT_DEBUG
	assert(debug);
#else
	assert(!debug);
#endif
	C->cb_root = 0;
	C->cb_ops = ops;
}

void
critbit_destroy(struct critbit *C)
{

	assert(C->cb_root == 0);
	assert(C->cb_ops != NULL);
	C->cb_ops = NULL;
}

void
critbit_node_init(struct critbit_node *N __dbgused __diagused,
    const void *key __diagused, size_t keylen __diagused)
{

	assert(((uintptr_t)N & 1) == 0);
	assert(((uintptr_t)&N->cbn_branch & 1) == 0);
	assert(key != NULL);
	assert(0 < keylen);
	assert(keylen <= CRITBIT_MAX_KEYLEN);

	DBG(N->cbn_magic = CRITBIT_NODE_MAGIC);
	DBG(N->cbn_branch.cbb_magic = 0); /* invalidate */
}

void
critbit_node_destroy(struct critbit_node *N __dbgused)
{

	DBG(assert(N->cbn_magic == CRITBIT_NODE_MAGIC));
	DBG(assert(N->cbn_branch.cbb_magic == 0));

	DBG(N->cbn_magic = 0);
}

static const uint8_t *
critbit_key(const struct critbit *C, const struct critbit_node *N,
    size_t *keylenp)
{
	const uint8_t *k;

	k = (*C->cb_ops->cbo_key)(C, N, keylenp);
	assert(k != NULL);
	assert(*keylenp <= CRITBIT_MAX_KEYLEN);

	return k;
}

/* Min/max and lookup */

struct critbit_node *
critbit_min(const struct critbit *C)
{
	struct critbit_branch *B;
	uintptr_t child;

	if (C->cb_root == 0)
		return NULL;

	for (child = C->cb_root; isbranch(child); child = B->cbb_child[0])
		B = child2branch(child);

	assert(isleaf(child));
	return child2leaf(child);
}

struct critbit_node *
critbit_max(const struct critbit *C)
{
	struct critbit_branch *B;
	uintptr_t child;

	if (C->cb_root == 0)
		return NULL;

	for (child = C->cb_root; isbranch(child); child = B->cbb_child[1])
		B = child2branch(child);

	assert(isleaf(child));
	return child2leaf(child);
}

struct critbit_node *
critbit_lookup(const struct critbit *C, const void *key, size_t keylen)
{
	struct critbit_node *N;
	size_t keylen0;

	assert(keylen <= CRITBIT_MAX_KEYLEN);

	/* Find the first node with the requested key as a prefix.  */
	N = critbit_lookup_prefix(C, key, keylen);
	if (N == NULL)
		return NULL;

	/* Reject if it is not an exact match.  */
	(void)critbit_key(C, N, &keylen0);
	if (keylen0 != keylen)
		return NULL;

	/* Success!  */
	return N;
}

/* Lookup prefix */

struct critbit_node *
critbit_lookup_prefix(const struct critbit *C, const void *key, size_t keylen)
{
	const uint8_t *const k = key;
	const struct critbit_branch *B;
	uintptr_t child;
	unsigned dir;
	struct critbit_node *N;
	const uint8_t *k0;
	size_t keylen0;
	size_t i;

	assert(keylen <= CRITBIT_MAX_KEYLEN);

	/* If the tree is empty, not there.  */
	if (C->cb_root == 0)
		return NULL;

	/* Find the first candidate match.  */
	for (child = C->cb_root; isbranch(child); child = B->cbb_child[dir]) {
		unsigned byte;

		B = child2branch(child);
		byte = (B->cbb_idx < keylen? k[B->cbb_idx] : 0);
		dir = (1 + (B->cbb_bits | byte)) >> 8;
	}

	/* Get the node.  */
	assert(isleaf(child));
	N = child2leaf(child);
	k0 = critbit_key(C, N, &keylen0);

	/* Reject if the prefix doesn't match.  */
	if (keylen0 < keylen)
		return NULL;
	for (i = 0; i < keylen; i++)
		if (k[i] != k0[i])
			return NULL;

	/* Success!  */
	return N;
}

/* Insert */

struct critbit_node *
critbit_insert(struct critbit *C, struct critbit_node *N)
{
	size_t keylen, keylen0;
	const uint8_t *k, *k0;
	struct critbit_branch *B;
	uintptr_t child, *childp;
	unsigned dir, byte = 0 /*XXXGCC*/, newdir, newbyte;
	struct critbit_node *N0;
	uint32_t i, diff, newbits;

	/* Get the key.  */
	k = critbit_key(C, N, &keylen);
	assert(k != NULL);
	assert(keylen <= CRITBIT_MAX_KEYLEN);

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

	/* Find the first candidate match.  */
	for (child = C->cb_root; isbranch(child); child = B->cbb_child[dir]) {
		B = child2branch(child);
		byte = (B->cbb_idx < keylen? k[B->cbb_idx] : 0);
		dir = (1 + (B->cbb_bits | byte)) >> 8;
	}

	/* Get the candidate node.  */
	assert(isleaf(child));
	N0 = child2leaf(child);
	k0 = critbit_key(C, N0, &keylen0);

	/* Find the index of the critical byte.  */
	for (i = 0; i < MIN(keylen, keylen0); i++) {
		newbyte = k[i];
		if ((diff = k0[i] ^ newbyte) != 0)
			goto branch;
	}
	if (keylen < keylen0) {		/* our key is shorter */
		newbyte = 0;
		do {
			if ((diff = k0[i]) != 0)
				break;
		} while (++i < keylen0);
		if (i == keylen0)	/* no 1 bits in k0, fake one at end */
			diff = 0x80;
		goto branch;
	} else if (keylen0 < keylen) {	/* their key is shorter */
		do {
			newbyte = k[i];
			if ((diff = newbyte) != 0)
				break;
		} while (++i < keylen);
		if (i == keylen)	/* no 1 bits in k, fake one at end */
			diff = newbyte = 0x80;
		goto branch;
	}

	return N0;		/* collision */

branch:
	/*
	 * Find the critical bit:
	 *
	 * . Set all bits from LSB up to the critical bit.
	 * . Clear all bits but the critical bit.
	 * . Complement -- set all bits but the critical bit.
	 *
	 * This way, newbits | b is 0xff iff b has 1 in the critical
	 * bit, so that (1 + (newbits | b)) >> 8 is 1 iff b has 1 in
	 * the critical bit.
	 */
	diff |= diff >> 1;
	diff |= diff >> 2;
	diff |= diff >> 4;
	diff ^= diff >> 1;
	newbits = diff ^ 0xff;

	/* Direction from parent in which to put the new item.  */
	newdir = (1 + (newbits | newbyte)) >> 8;

	/* Find where to insert a new branch.  */
	for (childp = &C->cb_root;
	     isbranch(*childp);
	     childp = &B->cbb_child[dir]) {
		B = child2branch(*childp);
		if (i < B->cbb_idx)
			break;
		if (i == B->cbb_idx && newbits < B->cbb_bits)
			break;
		byte = (B->cbb_idx < keylen? k[B->cbb_idx] : 0);
		dir = (1 + (B->cbb_bits | byte)) >> 8;
	}

	/* Initialize the new branch.  */
	B = &N->cbn_branch;
	B->cbb_idx = i;
	B->cbb_bits = newbits;
	B->cbb_child[newdir] = leaf2child(N);
	B->cbb_child[1 - newdir] = *childp;
	DBG(B->cbb_magic = CRITBIT_BRANCH_MAGIC);

	/* Insert it.  */
	*childp = branch2child(B);

	return N;		/* inserted */
}

/* Delete */

struct critbit_node *
critbit_delete_key(struct critbit *C, const void *key, size_t keylen)
{
	const uint8_t *const k = key;
	struct critbit_branch *B = NULL, *B0 = NULL;
	uintptr_t *childp, *parentp = NULL;
	unsigned dir = 0 /*XXXGCC*/, byte;
	struct critbit_node *N;
	const uint8_t *k0;
	size_t keylen0;
	uint32_t i;

	/* If the tree is empty, do nothing.  */
	if (C->cb_root == 0)
		return NULL;

	/* Find the closest node.  */
	for (childp = &C->cb_root;
	     isbranch(*childp);
	     childp = &B->cbb_child[dir]) {
		parentp = childp;
		B = child2branch(*childp);
		byte = (B->cbb_idx < keylen? k[B->cbb_idx] : 0);
		dir = (1 + (B->cbb_bits | byte)) >> 8;
	}

	/* Ignore it if it doesn't match.  */
	assert(isleaf(*childp));
	N = child2leaf(*childp);
	k0 = critbit_key(C, N, &keylen0);
	if (keylen0 != keylen)
		return NULL;
	for (i = 0; i < keylen; i++) {
		if (k0[i] != k[i])
			return NULL;
	}

	/* If the root was a leaf, just replace it by empty.  */
	if (parentp == NULL) {
		assert(childp == &C->cb_root);
		assert(C->cb_root == leaf2child(N));
		C->cb_root = 0;
		DBG(assert(N->cbn_branch.cbb_magic == 0));
		goto deleted_node;
	}

	/* Otherwise, replace it by the other child.  */
	assert(B != NULL);
	assert(B->cbb_child[dir] == leaf2child(N));
	*parentp = B->cbb_child[1 - dir];

	/* Is the branch we deleted N's branch?  If so, done!  */
	if (B == &N->cbn_branch) {
		DBG(assert(N->cbn_branch.cbb_magic == CRITBIT_BRANCH_MAGIC));
		goto deleted_branch;
	}

	/* Otherwise, find N's branch to replace it by the one we deleted.  */
	for (childp = &C->cb_root;
	     isbranch(*childp);
	     childp = &B0->cbb_child[dir]) {
		B0 = child2branch(*childp);
		if (B0 == &N->cbn_branch)
			break;
		byte = (B0->cbb_idx < keylen? k[B0->cbb_idx] : 0);
		dir = (1 + (B0->cbb_bits | byte)) >> 8;
	}
	assert(childp != NULL);

	/*
	 * If we didn't find it, then N is the distinguished node whose
	 * branch structure is not in use.  The branch we just deleted
	 * is now the distinguished one not in use.
	 */
	if (B0 != &N->cbn_branch) {
		DBG(assert(N->cbn_branch.cbb_magic == 0));
		DBG(assert(B->cbb_magic == CRITBIT_BRANCH_MAGIC));
		DBG(B->cbb_magic = 0);
		goto deleted_node;
	}

	/*
	 * If we found it, then the distinguished leaf node whose
	 * branch structure is not in use is not relevant.  Replace N's
	 * branch by the branch we just deleted.
	 */
	assert(B != &N->cbn_branch);
	assert(B != B0);
	assert(B0 == &N->cbn_branch);
	*B = *B0;
	*childp = branch2child(B);
	DBG(assert(B->cbb_magic == CRITBIT_BRANCH_MAGIC));
	DBG(assert(N->cbn_branch.cbb_magic == CRITBIT_BRANCH_MAGIC));

deleted_branch:
	/* We deleted N's branch.  */
	DBG(assert(N->cbn_branch.cbb_magic == CRITBIT_BRANCH_MAGIC));
	DBG(N->cbn_branch.cbb_magic = 0);
deleted_node:
	/* We deleted N.  */
	DBG(assert(N->cbn_branch.cbb_magic == 0));
	return N;
}

void
critbit_delete(struct critbit *C, struct critbit_node *N)
{
	struct critbit_node *N0;
	const void *key;
	size_t keylen;

	/* Get the key.  */
	key = critbit_key(C, N, &keylen);

	/* Delete by the key.  */
	N0 = critbit_delete_key(C, key, keylen);

	/* The node had better be the same one.  */
	assert(N0 == N);
}

/* Delete transaction -- not relevant for uniprocessor implementation */

void
critbit_tx_delete_begin(struct critbit *C __unused,
    struct critbit_tx_delete *td __unused,
    struct critbit_tx_delete_1 *td1 __unused, size_t n __unused)
{
}

struct critbit_node *
critbit_tx_delete_key(struct critbit *C, const void *key, size_t keylen,
    struct critbit_tx_delete *td __unused)
{

	return critbit_delete_key(C, key, keylen);
}

void
critbit_tx_delete(struct critbit *C, struct critbit_node *N,
    struct critbit_tx_delete *td __unused)
{

	critbit_delete(C, N);
}

void
critbit_tx_delete_commit(struct critbit *C __unused,
    struct critbit_tx_delete *td __unused)
{
}

/* Consistency check */

static void
critbit_check_leaf(const struct critbit *C, const struct critbit_node *N)
{
	const uint8_t *k;
	size_t keylen;
	const struct critbit_branch *B;
	uintptr_t child;
	unsigned dir;
	DBG(bool ok = false);

	k = critbit_key(C, N, &keylen);

	assert(C->cb_root != 0);
	for (child = C->cb_root; isbranch(child); child = B->cbb_child[dir]) {
		unsigned byte;

		B = child2branch(child);
		if (B == &N->cbn_branch)
			DBG(ok = true);
		byte = (B->cbb_idx < keylen? k[B->cbb_idx] : 0);
		dir = (1 + (B->cbb_bits | byte)) >> 8;
	}
	assert(isleaf(child));
	assert(N == child2leaf(child));
	DBG(assert(ok == (N->cbn_branch.cbb_magic != 0)));
}

static void
critbit_check_child(struct critbit *C, uintptr_t child, size_t idx __diagused)
{

	if (isbranch(child)) {
		struct critbit_branch *B = child2branch(child);

		assert(idx <= B->cbb_idx);

		critbit_check_child(C, B->cbb_child[0], B->cbb_idx);
		critbit_check_child(C, B->cbb_child[1], B->cbb_idx);
	} else {
		critbit_check_leaf(C, child2leaf(child));
	}
}

void
critbit_check(struct critbit *C)
{

	if (C->cb_root == 0)
		return;
	critbit_check_child(C, C->cb_root, 0);
}

/* Print */

/*
 * WARNING: This debugging operation requires that the node's key be a
 * NUL-terminated printable string.
 */

#include <stdio.h>
#include <strings.h>

static void
critbit_print_child(struct critbit *C, uintptr_t child, unsigned indent)
{

	if (isbranch(child)) {
		struct critbit_branch *const B = child2branch(child);

		printf("%*sbranch on 8*%u + %u (0x%02x) @ %p\n",
		    indent*2, "",
		    B->cbb_idx, 8 - ffs(0xff^B->cbb_bits), 0xff^B->cbb_bits,
		    B);

		critbit_print_child(C, B->cbb_child[0], indent + 1);
		critbit_print_child(C, B->cbb_child[1], indent + 1);
	} else {
		struct critbit_node *const N = child2leaf(child);
		size_t keylen __unused;
		const uint8_t *const k = critbit_key(C, N, &keylen);
		printf("%*sleaf @ %p: %s\n", indent*2, "", N, k);
	}
}

void
critbit_print(struct critbit *C)
{

	if (C->cb_root == 0) {
		printf("empty\n");
		return;
	}
	critbit_print_child(C, C->cb_root, 0);
}