/* 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));
}
