#include <fixed_rbtree_pool.h>
#include "os_macros.h"
/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

typedef fixed_rbtree_pool_t            rbtree_t;
typedef fixed_rbtree_pool_node_t       rbtree_node_t;


/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

#define RED     'R'
#define BLACK   'B'


/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


OS_STATIC_FORCE_INLINE
void node_init(rbtree_t* tree, rbtree_node_t* self, os_size_t obj_size, void* block, os_size_t block_size){
    self->left = self->right = self->parent = 0;
    self->color = RED;
    pool_init(&self->pool, obj_size);
    pool_add_block(&self->pool, block, block_size);
}

static
int node_cmp(const os_size_t obj_size, const rbtree_node_t* x, void* args){
    return (obj_size == x->pool.obj_size)?0:((obj_size > x->pool.obj_size)?1:-1);
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


OS_STATIC_FORCE_INLINE
void node_destroy(rbtree_t* tree, rbtree_node_t* x){
    pool_free(&tree->node_pool, x);
}

OS_STATIC_FORCE_INLINE
void rotate_left(rbtree_t* self, rbtree_node_t* x){
    rbtree_node_t * y = x->right;
    x->right = y->left;
    if(y->left){
        y->left->parent = x;
    }
    y->parent = x->parent;
    if(!x->parent){
        self->root = y;
    } else{
        if(x->parent->left==x){
            x->parent->left = y;
        }else{
            x->parent->right = y;
        }
    }
    y->left = x;
    x->parent = y;
}

OS_STATIC_FORCE_INLINE
void rotate_right(rbtree_t* self, rbtree_node_t* y){
    rbtree_node_t * x = y->left;
    y->left = x->right;
    if(x->right){
        x->right->parent = y;
    }
    x->parent = y->parent;
    if(!y->parent){
        self->root = x;
    }else{
        if(y==y->parent->left){
            y->parent->left = x;
        }else{
            y->parent->right = x;
        }
    }
    x->right = y;
    y->parent = x;
}

static rbtree_node_t* find(rbtree_t* self, const os_size_t obj_size, rbtree_node_t* x){
    while(x){
        int cmp = node_cmp(obj_size, x, 0);
        if(cmp < 0){
            x = x->left;
        }else if(cmp > 0){
            x = x->right;
        }else{
            return x;
        }
    }
    return NULL;
}


OS_STATIC_FORCE_INLINE
rbtree_node_t * node_parent(rbtree_node_t* x){
    return (x!=NULL)?x->parent:NULL;
}

OS_STATIC_FORCE_INLINE
char node_color(rbtree_node_t* x){
    return (x!=NULL)?x->color:BLACK;
}

OS_STATIC_FORCE_INLINE
bool node_is_red(rbtree_node_t* x){
    return (node_color(x)==RED)?true:false;
}

OS_STATIC_FORCE_INLINE
bool node_is_black(rbtree_node_t* x){
    return !node_is_red(x);
}

OS_STATIC_FORCE_INLINE
void node_set_black(rbtree_node_t* x){
    if(x){
        x->color = BLACK;
    }
}

OS_STATIC_FORCE_INLINE
void node_set_red(rbtree_node_t* x){
    if(x){
        x->color = RED;
    }
}

OS_STATIC_FORCE_INLINE
void node_set_color(rbtree_node_t* x, char color){
    if(x){
        x->color = color;
    }
}

OS_STATIC_FORCE_INLINE
void node_set_parent(rbtree_node_t* x, rbtree_node_t* parent){
    if(x){
        x->parent = parent;
    }
}


OS_STATIC_FORCE_INLINE
void node_insert_fixup(rbtree_t* tree, rbtree_node_t* node){
    rbtree_node_t * parent=0;
    rbtree_node_t * gparent=0;

    while(((parent = node->parent)!=NULL) && node_is_red(parent)){
        gparent = node_parent(parent);
        if(parent==gparent->left){
            rbtree_node_t * pUncle = gparent->right;
            if((pUncle!=NULL) && node_is_red(pUncle)){
                // Case 1条件：叔叔节点是红色
                node_set_black(pUncle);
                node_set_black(parent);
                node_set_red(gparent);
                node = gparent;
                continue;
            }

            if(parent->right==node){
                // Case 3条件：叔叔是黑色，且当前节点是右孩子
                rbtree_node_t * pTmp;
                rotate_left(tree, parent);
                pTmp = parent;
                parent = node;
                node = pTmp;
            }

            // Case 2条件：叔叔是黑色，且当前节点是左孩子。
            node_set_black(parent);
            node_set_red(gparent);
            rotate_right(tree, gparent);
        }else{
            //若“z的父节点”是“z的祖父节点的右孩子”
            rbtree_node_t * pUncle = gparent->left;
            if((pUncle!=NULL) && node_is_red(pUncle)) {
                // Case 1条件：叔叔节点是红色
                node_set_black(pUncle);
                node_set_black(parent);
                node_set_red(gparent);
                node = gparent;
                continue;
            }

            // Case 2条件：叔叔是黑色，且当前节点是左孩子
            if(parent->left == node){
                rbtree_node_t * pTmp;
                rotate_right(tree, parent);
                pTmp = parent;
                parent = node;
                node = pTmp;
            }

            // Case 3条件：叔叔是黑色，且当前节点是右孩子。
            node_set_black(parent);
            node_set_red(gparent);
            rotate_left(tree, gparent);
        }
    }
    node_set_black(tree->root);
}

OS_STATIC_FORCE_INLINE
void node_insert(rbtree_t* tree, rbtree_node_t* node){
    int cmp = 0;
    rbtree_node_t * y = NULL;
    rbtree_node_t * x = tree->root;

    // 1. 将红黑树当作一颗二叉查找树，将节点添加到二叉查找树中。
    while (x) {
        y = x;
        cmp = node_cmp(node->pool.obj_size, x,0);
        if (cmp < 0)
            x = x->left;
        else
            x = x->right;
    }
    // y 是要插入的父节点
    node->parent = y;
    if(!y){
        // 插入根节点
        tree->root = node;
    }else{
        // 判断插入父节点的左边还是右边
        cmp = node_cmp(node->pool.obj_size, y, 0);
        if(cmp < 0){
            y->left = node;
        }else{
            y->right = node;
        }
    }

    // 2. 设置节点的颜色为红色
    node->color = RED;

    // 3. 将它重新修正为一颗二叉查找树
    node_insert_fixup(tree, node);
}

OS_STATIC_FORCE_INLINE
void node_remove_fixeup(rbtree_t* tree, rbtree_node_t* node, rbtree_node_t* parent){
    rbtree_node_t * other;

    while(parent!=NULL && (node==NULL || node_is_black(node)) && (node!=tree->root) ){
        if(parent->left == node){
            other= parent->right;
            if(node_is_red(other)){
                // Case 1: x的兄弟w是红色的
                node_set_black(other);
                node_set_red(parent);
                rotate_left(tree, parent);
                other = parent->right;
            }

            if((other->left==NULL || node_is_black(other->left)) && (other->right==NULL ||
                                                                     node_is_black(other->right))){
                // Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的
                node_set_red(other);
                node = parent;
                parent = node_parent(node);
            }else{
                if(other->right==NULL || node_is_black(other->right)){
                    // Case 4: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。
                    node_set_black(other->left);
                    node_set_red(other);
                    rotate_right(tree, other);
                    other = parent->right;
                }
                // Case 3: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。
                node_set_color(other, node_color(parent));
                node_set_black(parent);
                node_set_black(other->right);
                rotate_left(tree, parent);
                node = tree->root;
                break;
            }
        }else{
            other= parent->left;
            if(node_is_red(other)){
                // Case 1: x的兄弟w是红色的
                node_set_black(other);
                node_set_red(parent);
                rotate_right(tree, parent);
                other = parent->left;
            }

            if((other->left==NULL || node_is_black(other->left)) && (other->right==NULL ||
                                                                     node_is_black(other->right))){
                // Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的
                node_set_red(other);
                node = parent;
                parent = node_parent(node);
            }else{
                if(other->left==NULL || node_is_black(other->left)){
                    // Case 4: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。
                    node_set_black(other->right);
                    node_set_red(other);
                    rotate_left(tree, other);
                    other= parent->left;
                }

                // Case 3: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。
                node_set_color(other, node_color(parent));
                node_set_black(parent);
                node_set_black(other->left);
                rotate_right(tree, parent);
                node = tree->root;
                break;
            }
        }
    }
    if(node){
        node_set_black(node);
    }
}

OS_STATIC_FORCE_INLINE
void node_remove(rbtree_t* tree, rbtree_node_t* node){
    rbtree_node_t * child=NULL;
    rbtree_node_t * parent=NULL;
    char color;

    // 被删除节点的"左右孩子都不为空"的情况。
    if((node->left!=NULL) && (node->right!=NULL)){
        // 被删节点的后继节点。(称为"取代节点")
        // 用它来取代"被删节点"的位置，然后再将"被删节点"去掉。

        rbtree_node_t * replace = node;

        // 获取后继节点
        replace = replace->right;
        while(replace->left){
            replace = replace->left;
        }

        // "node节点"不是根节点(只有根节点不存在父节点)
        if(node_parent(node)!=NULL){
            if(node->parent->left == node){
                node->parent->left = replace;
            }else{
                node->parent->right = replace;
            }
        }else{
            // "node节点"是根节点，更新根节点。
            tree->root = replace;
        }

        // child是"取代节点"的右孩子，也是需要"调整的节点"。
        // "取代节点"肯定不存在左孩子！因为它是一个后继节点。

        child = replace->right;
        parent= node_parent(replace);

        // 保存"取代节点"的颜色
        color = node_color(replace);

        // "被删除节点"是"它的后继节点的父节点"
        if(parent == node){
            parent = replace;
        }else {
            // child 不为空
            if(child!=NULL){
                node_set_parent(child, parent);
            }
            parent->left = child;

            replace->right = node->right;
            node_set_parent(node->right, replace);
        }

        replace->parent = node->parent;
        replace->color = node->color;
        replace->left= node->left;
        node->left->parent = replace;

        if(color==BLACK){
            node_remove_fixeup(tree, child, parent);
        }

        return;
    }

    if(node->left!=NULL){
        child = node->left;
    }else{
        child= node->right;
    }

    parent = node->parent;
    color = node->color;
    if(child!=NULL){
        child->parent = parent;
    }

    // "node节点"不是根节点
    if(parent){
        if(parent->left==node){
            parent->left = child;
        }else{
            parent->right = child;
        }
    }else{
        tree->root = child;
    }


    if(color==BLACK){
        node_remove_fixeup(tree, child, parent);
    }

}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


void fixed_rbtree_pool_init(fixed_rbtree_pool_t* self, void* node_block, os_size_t node_block_size){
    self->root = 0;
    pool_init(&self->node_pool, sizeof(fixed_rbtree_pool_node_t));
    pool_add_block(&self->node_pool, node_block, node_block_size);
}

void fixed_rbtree_pool_add_node_pool(fixed_rbtree_pool_t* self, void* node_block, os_size_t node_block_size){
    pool_add_block(&self->node_pool, node_block, node_block_size);
}

void fixed_rbtree_pool_destroy(fixed_rbtree_pool_t* self){
    while(self->root){
        rbtree_node_t * p = self->root;
        node_remove(self, p);
        node_destroy(self, p);
    }
}

void fixed_rbtree_pool_add(fixed_rbtree_pool_t* self, os_size_t obj_size, void* block, os_size_t block_size){
    rbtree_node_t* x = find(self, obj_size, self->root);
    if(x==NULL){
        x = pool_alloc(&self->node_pool);
        node_init(self, x, obj_size, block, block_size);
        node_insert(self, x);
    }else{
        pool_add_block(&x->pool, block, block_size);
    }
}

void* fixed_rbtree_pool_alloc(fixed_rbtree_pool_t* self, os_size_t obj_size){
    rbtree_node_t * x = find(self, obj_size, self->root);
    OS_ASSERT(x, "No pool exist for obj_size: %"OS_PRIu, obj_size);
    if(x==NULL){
        return NULL;
    }
    return pool_alloc(&x->pool);
}

void fixed_rbtree_pool_free(fixed_rbtree_pool_t* self, os_size_t obj_size, void* ptr){
    if(!ptr) return;
    rbtree_node_t * x = find(self, obj_size, self->root);
    OS_ASSERT(x, "No pool exist for obj_size: %"OS_PRIu, obj_size);
    if(x==NULL){
        return;
    }
    pool_free(&self->node_pool, ptr);
}