RTOS/DeviceKit/drv_rbtree.c

#include <drv_rbtree.h>

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

typedef drv_rbtree_node_t   rbtree_node_t;
typedef drv_rbtree_t        rbtree_t;

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

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

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

OS_STATIC_FORCE_INLINE
int node_cmp(const char* a, const char * b){
    os_size_t na_len = strlen(a);
    os_size_t nb_len = strlen(b);
    return strncmp(a, b, OS_MIN(na_len, nb_len));
}

OS_STATIC_FORCE_INLINE
void node_init(rbtree_t* tree, rbtree_node_t * self, drv_device_t* val){
    self->left = self->right = self->parent = 0;
    self->color = RED;
    self->device = val;
}

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;
}

OS_STATIC_FORCE_INLINE
rbtree_node_t * find(rbtree_t* self, const char* key, rbtree_node_t * x){
    while(x){
        int cmp = node_cmp(key, x->device->name);
        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;
    }
}


static
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);
}

static
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->device->name, x->device->name);
        if (cmp < 0)
            x = x->left;
        else
            x = x->right;
    }
    // y 是要插入的父节点
    node->parent = y;
    if(!y){
        // 插入根节点
        tree->root = node;
    }else{
        // 判断插入父节点的左边还是右边
        cmp = node_cmp(node->device->name, y->device->name);
        if(cmp < 0){
            y->left = node;
        }else{
            y->right = node;
        }
    }

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

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

static
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);
    }
}

static
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 drv_rbtree_init(drv_rbtree_t* self, void* node_block, os_size_t node_block_size){
    pool_init(&self->node_pool, sizeof(rbtree_node_t));
    pool_add_block(&self->node_pool, node_block, node_block_size);
    self->root = 0;
}

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

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

drv_rbtree_node_t* drv_rbtree_find(drv_rbtree_t* self, const char* key){
    if(key==NULL) return NULL;
    return find(self, key, self->root);
}

drv_rbtree_node_t* drv_rbtree_custom_find(drv_rbtree_t* self, const char* key
        , drv_rbtree_node_t* (*custom_find)(rbtree_t* tree, const char* key, drv_rbtree_node_t* x, void* args)
        , void* args)
{
    if(key==NULL) return NULL;
    return custom_find(self, key, self->root, args);
}

os_err_t drv_rbtree_insert(drv_rbtree_t* self, drv_device_t* device){
    drv_rbtree_node_t* p = pool_alloc(&self->node_pool);
    if(!p){
        return OS_ERR_FAIL;
    }
    node_init(self, p, device);
    node_insert(self, p);
    return OS_ERR_OK;
}

void drv_rbtree_remove(drv_rbtree_t* self, const char* key){
    drv_rbtree_node_t * pNode = drv_rbtree_find(self, key);
    if(pNode==NULL){
        return;
    }
    node_remove(self, pNode);
    node_destroy(self, pNode);
}

void drv_rbtree_remove_node(drv_rbtree_t* self, drv_rbtree_node_t* x){
    if(x==NULL) return;
    node_remove(self, x);
    node_destroy(self, x);
}

static void in_order(drv_rbtree_t* self, drv_rbtree_node_t* x,
                     int (*apply)(drv_rbtree_t* tree, drv_rbtree_node_t* x, void* args), void* args){
    if(x){
        in_order(self, x->left, apply, args);
        if(apply(self, x, args)==0){
            return;
        }
        in_order(self, x->right, apply, args);
    }
}

void drv_rbtree_inorder(drv_rbtree_t* self
        , int (*apply)(drv_rbtree_t* tree, drv_rbtree_node_t* x, void* args), void* args){
    in_order(self, self->root, apply, args);
}