/*
 * CS5600 University of Utah
 * Charles McGarvey
 * mcgarvey@eng.utah.edu
 */

#include "rbtree.h"


#define BLACK       (RBTREE_BLACK)
#define RED         (RBTREE_RED)
#define SENTINEL    (rbtree_sentinel)


rbtree_node_t _rbtree_sentinel = {
    &_rbtree_sentinel, &_rbtree_sentinel, &_rbtree_sentinel,
    BLACK
};
rbtree_node_t* rbtree_sentinel = &_rbtree_sentinel;


/*
 * Allocate and initialize a new node for the tree.
 */
INLINE_MAYBE
rbtree_node_t* rbtree_node_alloc(rbtree_t* t)
{
    rbtree_node_t* n = mem_alloc(sizeof(rbtree_node_t) + t->size);
    n->parent = NULL;
    n->left = n->right = rbtree_sentinel;
    n->color = RBTREE_BLACK;
    return n;
}


INLINE_MAYBE
void rbtree_init(rbtree_t* t, size_t size, rbtree_search_fn_t search)
{
    t->root = rbtree_sentinel;
    t->size = size;
    t->search = search;
}

rbtree_t* rbtree_alloc(size_t size, rbtree_search_fn_t search)
{
    rbtree_t* t = mem_alloc(sizeof(rbtree_t));
    rbtree_init(t, size, search);
    return t;
}


void rbtree_destroy(rbtree_t* t)
{
    rbtree_clear(t);
    mem_free(t);
}


INLINE_MAYBE
void _rotate_left(rbtree_node_t** r, rbtree_node_t* n)
{
    rbtree_node_t* t = n->right;
    n->right = t->left;

    if (t->left != SENTINEL) {
        t->left->parent = n;
    }
    t->parent = n->parent;

    if (n == *r) {
        *r = t;
    }
    else if (n == n->parent->left) {
        n->parent->left = t;
    }
    else {
        n->parent->right = t;
    }

    t->left = n;
    n->parent = t;
}

INLINE_MAYBE
void _rotate_right(rbtree_node_t** r, rbtree_node_t* n)
{
    rbtree_node_t* t = n->left;
    n->left = t->right;

    if (t->right != SENTINEL) {
        t->right->parent = n;
    }
    t->parent = n->parent;

    if (n == *r) {
        *r = t;
    }
    else if (n == n->parent->right) {
        n->parent->right = t;
    }
    else {
        n->parent->left = t;
    }

    t->right = n;
    n->parent = t;
}


rbtree_node_t* rbtree_insert(rbtree_t* t, void* data)
{
    rbtree_node_t* r = t->root;
    rbtree_node_t* n = r;

    int c = t->search(&n, data);
    if (c == 0) {
        memcpy(RBTREE_NODE_DATA(n), data, t->size);
        return n;
    }
    else if (n == SENTINEL) {
        r = rbtree_node_alloc(t);
        memcpy(RBTREE_NODE_DATA(r), data, t->size);
        t->root = r;
        return r;
    }

    rbtree_node_t* m = rbtree_node_alloc(t);
    m->parent = n;
    m->color = RED;
    memcpy(RBTREE_NODE_DATA(m), data, t->size);
    if (c < 0) {
        n->left = m;
    }
    else {
        n->right = m;
    }

    n = m;
    while (n != r && n->parent->color == RED) {
        if (n->parent == n->parent->parent->left) {
            rbtree_node_t* t = n->parent->parent->right;
            if (t->color == RED) {
                n->parent->color = BLACK;
                t->color = BLACK;
                n->parent->parent->color = RED;
                n = n->parent->parent;
            }
            else {
                if (n == n->parent->right) {
                    n = n->parent;
                    _rotate_left(&r, n);
                }
                n->parent->color = BLACK;
                n->parent->parent->color = RED;
                _rotate_right(&r, n->parent->parent);
            }
        }
        else {
            rbtree_node_t* t = n->parent->parent->left;
            if (t->color == RED) {
                n->parent->color = BLACK;
                t->color = BLACK;
                n->parent->parent->color = RED;
                n = n->parent->parent;
            }
            else {
                if (n == n->parent->left) {
                    n = n->parent;
                    _rotate_right(&r, n);
                }
                n->parent->color = BLACK;
                n->parent->parent->color = RED;
                _rotate_left(&r, n->parent->parent);
            }
        }
    }
    t->root = r;
    r->color = BLACK;

    return m;
}

rbtree_node_t* rbtree_delete(rbtree_t* t, void* data)
{
    rbtree_node_t* n = t->root;
    if (t->search(&n, data) == 0) {
        return rbtree_delete_node(t, n);
    }
    return NULL;
}

rbtree_node_t* rbtree_search(rbtree_t* t, void* data)
{
    rbtree_node_t* n = t->root;
    if (t->search(&n, data) == 0) {
        return n;
    }
    return NULL;
}


rbtree_node_t* rbtree_delete_node(rbtree_t* t, rbtree_node_t* n)
{
    rbtree_node_t* r = t->root;
    rbtree_node_t* u;
    rbtree_node_t* v;

    if (n->left == SENTINEL) {
        u = n->right;
        v = n;
    }
    else if (n->right == SENTINEL) {
        u = n->left;
        v = n;
    }
    else {
        v = rbtree_node_min(n->right);
        if (v->left != SENTINEL) {
            u = v->left;
        }
        else {
            u = v->right;
        }
    }

    if (v == r) {
        r = u;
        u->color = BLACK;
        goto done;
    }

    int isred = (v->color == RED);

    if (v == v->parent->left) {
        v->parent->left = u;
    }
    else {
        v->parent->right = u;
    }

    if (v == n) {
        u->parent = v->parent;
    }
    else {
        if (v->parent == n) {
            u->parent = v;
        }
        else {
            u->parent = v->parent;
        }

        *v = *n;

        if (n == r) {
            r = v;
        }
        else if (n == n->parent->left) {
            n->parent->left = v;
        }
        else {
            n->parent->right = v;
        }

        if (v->left != SENTINEL) {
            v->left->parent = v;
        }
        if (v->right != SENTINEL) {
            v->right->parent = v;
        }
    }

    if (isred) {
        goto done;
    }

    while (u != r && u->color == BLACK) {
        if (u == u->parent->left) {
            rbtree_node_t* w = u->parent->right;
            if (w->color == RED) {
                w->color = BLACK;
                w->parent->color = RED;
                _rotate_left(&r, u->parent);
                w = u->parent->right;
            }
            if (w->left->color == BLACK && w->right->color == BLACK) {
                w->color = RED;
                u = u->parent;
            }
            else {
                if (w->right->color == BLACK) {
                    w->left->color = BLACK;
                    w->color = RED;
                    _rotate_right(&r, w);
                    w = u->parent->right;
                }
                w->color = u->parent->color;
                u->parent->color = BLACK;
                w->right->color = BLACK;
                _rotate_left(&r, u->parent);
                u = r;
            }
        }
        else {
            rbtree_node_t* w = u->parent->left;
            if (w->color == RED) {
                w->color = BLACK;
                w->parent->color = RED;
                _rotate_right(&r, u->parent);
                w = u->parent->left;
            }
            if (w->left->color == BLACK && w->right->color == BLACK) {
                w->color = RED;
                u = u->parent;
            }
            else {
                if (w->left->color == BLACK) {
                    w->right->color = BLACK;
                    w->color = RED;
                    _rotate_left(&r, w);
                    w = u->parent->left;
                }
                w->color = u->parent->color;
                u->parent->color = BLACK;
                w->left->color = BLACK;
                _rotate_right(&r, u->parent);
                u = r;
            }
        }
    }
    
done:
    t->root = r;
    n->parent = n->left = n->right = NULL;
    return n;
}


void node_call(rbtree_node_t* n, rbtree_call_fn_t fn)
{
    if (n == SENTINEL) {
        return;
    }
    node_call(n->left, fn);
    fn(RBTREE_NODE_DATA(n));
    node_call(n->right, fn);
}

void rbtree_call(const rbtree_t* t, rbtree_call_fn_t fn)
{
    node_call(t->root, fn);
}


/*
 * Prune off a branch of the tree.
 */
static void _node_prune(rbtree_node_t* n)
{
    if (n == SENTINEL) {
        return;
    }
    _node_prune(n->left);
    _node_prune(n->right);
    mem_free(n);
}

void rbtree_clear(rbtree_t* t)
{
    if (t == NULL) {
        return;
    }
    _node_prune(t->root);
}