/* djb's critbit with associated data storage.
 *
 * Based on https://github.com/agl/critbit, which is in public domain.
 * Changes:
 *   - data storage added
 *   - cweb removed
 *   - functions renamed and data types cleaned to my liking
 */
#include <sys/types.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>

#include "cbtree.h"

struct node {
  void *child[2];
  uint32_t byte;
  uint8_t otherbits;
};

/* our own memory management */

struct pool
{
	struct pool *next;
	size_t pool_size;
	size_t free_index;
	char mem[0];
};

static struct pool *internal = NULL;
static struct pool *external = NULL;

static int new_pool(struct pool **root, size_t size)
{
	struct pool *pool;

	size = (size + sizeof(void *) - 1) / sizeof(void *);
	size *= sizeof(void *);
	pool = malloc(size + sizeof(struct pool));
	if (!pool) {
		return 1;
	}
	pool->next = *root;
	pool->free_index = 0;
	pool->pool_size = size;
	*root = pool;
	return 0;
}

static void *alloc(struct pool **root, size_t size)
{
	void *ret;
	size = (size + sizeof(void *) - 1) / sizeof(void *);
	size *= sizeof(void *);
	if (!*root)
		if (new_pool(root, size > 256000 ? size : 256000) != 0)
			return NULL;
	if ((*root)->pool_size - (*root)->free_index < size)
		if (new_pool(root, size > 256000 ? size : 256000) != 0)
			return NULL;
	ret = (*root)->mem + (*root)->free_index;
	(*root)->free_index += size;
	return ret;
}

/* main code */

intptr_t*
cbtree_find(struct cbtree *t, const char *u)
{
    const uint8_t *ubytes = (void *)u;
    const size_t ulen = strlen(u);
    uint8_t *p = t->root;

	/* Test for empty tree */
	if (!p) return NULL;

    /* Walk tree for best member */
	while (1 & (intptr_t) p) {
		struct node *q = (void *)(p - 1);
		/* Calculate direction */
		int direction;
		uint8_t c = 0;

		if (q->byte < ulen)
			c = ubytes[q->byte];
		direction = (1 + (q->otherbits | c)) >> 8;

		p = q->child[direction];
	}

	/* The leaves contain "[data ptr][string]" */
	if (strcmp(u, (const char *)(p + sizeof(intptr_t))) == 0)
		return (intptr_t *)p;
	return NULL;
}

intptr_t*
cbtree_insert(struct cbtree *t, const char *u)
{
	const uint8_t *const ubytes = (void *) u;
	const size_t ulen = strlen(u);
	uint8_t *p = t->root;

	/* Deal with inserting into an empty tree */
	if (!p) {
		char *x = alloc(&external, ulen + 1 + sizeof(intptr_t));
		if (!x)
			return NULL;
		*((intptr_t *)x) = 0;
		memcpy(x + sizeof(intptr_t), u, ulen + 1);
		t->root = x;
		return (intptr_t *)x;
	}

    /* Walk tree for best member */
	while (1 & (intptr_t) p) {
		struct node *q = (void *)(p - 1);
		/* Calculate direction */
		int direction;
		uint8_t c = 0;

		if (q->byte < ulen)
			c = ubytes[q->byte];
		direction = (1 + (q->otherbits | c)) >> 8;

		p = q->child[direction];
	}

	/* Find the critical bit */
	/* 1: Find differing byte */
	uint32_t newbyte;
	uint32_t newotherbits;

	for (newbyte = 0; newbyte < ulen; ++newbyte) {
		if (p[sizeof(intptr_t) + newbyte] != ubytes[newbyte]) {
			newotherbits = p[sizeof(intptr_t) + newbyte] ^ ubytes[newbyte];
			goto different_byte_found;
		}
	}

	if (p[sizeof(intptr_t) + newbyte] != 0) {
		newotherbits = p[sizeof(intptr_t) + newbyte];
		goto different_byte_found;
	}
	return (intptr_t *)p;

different_byte_found:

	/* 2: Find differing bit */
	newotherbits |= newotherbits >> 1;
	newotherbits |= newotherbits >> 2;
	newotherbits |= newotherbits >> 4;
	newotherbits = (newotherbits & ~(newotherbits >> 1)) ^ 255;
	uint8_t c = p[sizeof(intptr_t) + newbyte];
	int newdirection = (1 + (newotherbits | c)) >> 8;

	/* Insert new string */

	/* 1: Allocate new node structure */
	struct node *newnode;

	newnode = alloc(&internal, sizeof(struct node));
	if (!newnode)
		return NULL;

	char *x = alloc(&external, ulen + 1 + sizeof(intptr_t));
	if (!x)
		return NULL;
	*((intptr_t *)x) = 0;
	memcpy(x + sizeof(intptr_t), ubytes, ulen + 1);

	newnode->byte = newbyte;
	newnode->otherbits = newotherbits;
	newnode->child[1 - newdirection] = x;

	/* 2: Insert new node */
	void **wherep = &t->root;
	for (;;) {
		uint8_t *p = *wherep;
		if (!(1 & (intptr_t) p)) break;
		struct node *q = (void *) (p - 1);
		if (q->byte > newbyte) break;
		if (q->byte == newbyte && q->otherbits > newotherbits) break;
		uint8_t c = 0;
		if (q->byte < ulen) c = ubytes[q->byte];
		const int direction = (1 + (q->otherbits | c)) >> 8;
		wherep = q->child + direction;
	}

	newnode->child[newdirection] = *wherep;
	*wherep = (void *) (1 + (char *) newnode);

	return (intptr_t *)x;
}

intptr_t
cbtree_delete(struct cbtree *t, const char *u)
{
	const uint8_t *ubytes = (void *) u;
	const size_t ulen = strlen(u);
	uint8_t *p = t->root;
	void **wherep = &t->root;
	void **whereq = 0;
	struct node *q = 0;
	int direction = 0;
	intptr_t ret;

	/* Deal with deleting from an empty tree */
	if (!p) return 0;

	/* Walk the tree for the best match */
	while (1 & (intptr_t) p) {
		whereq = wherep;
		q = (void *) (p - 1);
		uint8_t c = 0;
		if (q->byte < ulen) c = ubytes[q->byte];
		direction = (1 + (q->otherbits | c)) >> 8;
		wherep = q->child + direction;
		p = *wherep;
	}

	/* Check the best match */
	if (0 != strcmp(u, (const char *)(p + sizeof(intptr_t))))
		return 0;
	ret = *((intptr_t *)p);

	/* Remove the element and/or node */
	if (!whereq) {
		t->root = 0;
		return ret;
	}

	*whereq = q->child[1 - direction];
	// free(q);

	return ret;
}

static void
traverse(void *top) {
	uint8_t *p = top;

	if (1 & (intptr_t) p) {
		struct node *q = (void *) (p - 1);
		traverse(q->child[0]);
		traverse(q->child[1]);
		// free(q);
	} else {
		// free(p);
	}
}

void
cbtree_clear(struct cbtree *t)
{
	if (t->root) traverse(t->root);
	t->root = NULL;
}

static int
allprefixed_traverse(uint8_t *top, int (*handle)(const char *, intptr_t *, void *), void *arg)
{
	int direction;

	/* Deal with an internal node */
	if (1 & (intptr_t) top) {
		struct node *q = (void *) (top - 1);
		for (direction = 0; direction < 2; ++direction)
			switch(allprefixed_traverse(q->child[direction], handle, arg)) {
			case 1: break;
			case 0: return 0;
			default: return -1;
			}
		return 1;
	}

	/* Deal with an external node */
	return handle((const char *)(top + sizeof(intptr_t)), (intptr_t *)top, arg);
}

int
cbtree_allprefixed(struct cbtree *t, const char *prefix,
				   int (*handle)(const char *, intptr_t *, void *),
				   void *arg)
{
	const uint8_t *ubytes = (void *) prefix;
	const size_t ulen = strlen(prefix);
	uint8_t *p = t->root;
	uint8_t *top = p;
	int i;

	if (!p) return 1; /* S = $\emptyset$ */

	/* Walk tree, maintaining top pointer */
	while (1 & (intptr_t) p) {
		struct node *q = (void *) (p - 1);
		uint8_t c = 0;
		if (q->byte < ulen) c = ubytes[q->byte];
		const int direction = (1 + (q->otherbits | c)) >> 8;
		p = q->child[direction];
		if (q->byte < ulen) top = p;
	}

	/* Check prefix */
	for (i = 0; i < ulen; ++i) {
		if (p[i+sizeof(intptr_t)] != ubytes[i]) return 1;
	}

	return allprefixed_traverse(top, handle, arg);
}

static const char *byte_to_binary(int x)
{
	static char b[9];
	int z;

	b[0] = '\0';
	for (z = 128; z > 0; z >>= 1)
		strcat(b, ((x & z) == z) ? "1" : "0");

	return b;
}

static void
traverse_dump(void *top, int level, int byte)
{
	uint8_t *p = top;
	int i;

	for (i = 0; i < level; i++) printf(" ");
	if (1 & (intptr_t) p) {
		struct node *q = (void *) (p - 1);
		printf("[byte(%d),otherbits(%s)]\n", q->byte, byte_to_binary(q->otherbits));
		traverse_dump(q->child[0], level + 1, q->byte);
		traverse_dump(q->child[1], level + 1, q->byte);
	} else {
		const size_t ulen = strlen((char *)(p + sizeof(intptr_t)));
		int c = byte < ulen ? p[sizeof(intptr_t) + byte] : 0;
		printf("\"%s\" (%s)\n", p + sizeof(intptr_t), byte_to_binary(c));
	}
}

void
cbtree_dump(struct cbtree *t)
{
	if (t->root)
		traverse_dump(t->root, 0, 0);
	printf("\n");
}

char*
cbtree_next(struct cbtree *t, const char *u, intptr_t *data)
{
	const uint8_t *ubytes = (void *) u;
	const size_t ulen = strlen(u);
	uint8_t *p = t->root;
	uint8_t *branch = NULL;

	if (!p) return NULL;

	/* Walk tree, maintaining top pointer */
	while (1 & (intptr_t) p) {
		struct node *q = (void *) (p - 1);
		uint8_t c = 0;
		if (q->byte < ulen) c = ubytes[q->byte];
		const int direction = (1 + (q->otherbits | c)) >> 8;
		if (direction == 0)
			branch = q->child[1];
		p = q->child[direction];
	}

	/* check whether what we found is what we are looking for already */
	if (strcmp((char *)(p + sizeof(intptr_t)), u) > 0) {
		if (data) *data = *((intptr_t *)p);
		return (char *)(p + sizeof(intptr_t));
	}

	if (!branch) return NULL;

	/* select the lowest value on the branch */
	p = branch;
	while (1 & (intptr_t) p) {
		struct node *q = (void *) (p - 1);
		p = q->child[0];
	}
	if (data) *data = *((intptr_t *)p);
	return (char *)(p + sizeof(intptr_t));
}