linux内核红黑树运用小实例

linux内核版本linux-3.10.36
在linux内核源码中，红黑树是一个比较独立的模块，很容易将其剥离出来，拿到应用层使用。

结构

linux内核的rb_node结构体

struct rb_node {
    unsigned long  __rb_parent_color;
    struct rb_node *rb_right;
    struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));

rb_parent_color这个域其实同时包含了颜色信息以及父亲节点的指针，因为该域是一个long的类型，需要大小为sizeof(long)的对齐，那么在一般的32位机器上，其后两位的数值永远是0，于是可以拿其中的一位来表示颜色。事实上，这里就是使用了最低位来表示颜色信息

根节点指针

struct rb_root {
    struct rb_node *rb_node;
};

接口

可以通过调用rb_replace_node来替换一个节点，但是替换完成后并不会对红黑树做任何调整，所以如果新节点的值与被替换的值有所不同时，可能会出现问题。

void rb_replace_node(struct rb_node *old, struct rb_node *new, struct rb_root *tree);

插入与删除

extern void rb_insert_color(struct rb_node *, struct rb_root *);//对新插入的节点着色，并修正红黑树使其达到平衡
extern void rb_erase(struct rb_node *, struct rb_root *);

插入节点时需要把新节点指向其父亲节点,通过该函数完成

static inline void rb_link_node(struct rb_node * node, struct rb_node * parent,
                struct rb_node ** rb_link)

rb_entry仅仅是container_of的封装而已

#define rb_entry(ptr, type, member) container_of(ptr, type, member)

其它接口在下面的代码中有所解释

例子

arvik写了一个以mac为索引的红黑树小例子，里面包括增，删，查，销毁几个操作，供大家参考。如下：
头文件

/****************************
* author:           arvik
* blog:             http://blog.csdn.net/u012819339
* email:            1216601195@qq.com
* kernel: linux-3.10.36
****************************/
#include <linux/string.h>
#include <linux/rbtree.h>
#include <linux/types.h>  
#include <linux/slab.h>

#define IPV4_MAC_LEN    6
typedef struct user_rb_s
{
    struct rb_node user_node;

    uint8_t umac[IPV4_MAC_LEN];
}user_rb_t;

#define mac_cmp(smac, dmac)     memcmp(smac, dmac, IPV4_MAC_LEN)

/*
申请内存存放节点
*/
user_rb_t *create_user_node(uint8_t *mac)
{
    user_rb_t *node;
    node = (user_rb_t *)kmalloc(sizeof(user_rb_t), GFP_ATOMIC);
    if(node == NULL)
        return node;

    memset(node, 0, sizeof(user_rb_t));
    memcpy(node->umac, mac, IPV4_MAC_LEN);
    return node;
}

/*
描述：查找节点
返回：节点指针
*/
user_rb_t   *user_rb_find(struct rb_root *root, uint8_t *umac)
{
    struct rb_node *node = root->rb_node;
    user_rb_t *goal_node = NULL;
    int res=0;

    while(node)
    {
        goal_node = rb_entry(node, user_rb_t, user_node);
        res = mac_cmp(umac, goal_node->umac);
        if(res<0)
            node = node->rb_left;
        else if(res>0)
            node = node->rb_right;
        else
            return goal_node;
    }

    return NULL;
}


//rb_insert
/*
描述：插入节点
返回：true or FLASE
*/
int user_rb_insert(struct rb_root *root, uint8_t *umac)
{
    struct rb_node **new = &(root->rb_node);
    struct rb_node *parent = NULL;
    user_rb_t *goal_node = NULL;
    int res=0;

    while(*new)
    {
        parent = *new;
        goal_node = rb_entry(*new, user_rb_t, user_node);
        res = mac_cmp(umac, goal_node->umac);
        if(res<0)
            new = &((*new)->rb_left);
        else if(res>0)
            new = &((*new)->rb_right);
        else
            return true;    
    }

    goal_node = create_user_node(umac);
    if(goal_node == NULL)
        return false;

    //add new node
    rb_link_node(&goal_node->user_node, parent, new);
    //rebalance rbtree
    rb_insert_color(&goal_node->user_node, root);

    return true;
}

/*
描述：删除单个节点
*/
void user_rb_delete(struct rb_root *root, uint8_t *umac)
{
    user_rb_t *goal_node = user_rb_find(root, umac);
    if(goal_node != NULL)
    {
        rb_erase(&goal_node->user_node, root);
        kfree(goal_node);
    }
}

/*
描述：删除整颗红黑树
*/
void user_rb_destory(struct rb_root *root)
{
    struct rb_node *node = NULL, *tmp_node = NULL;
    user_rb_t *goal_node = NULL;

    for(node = rb_first(root); node; )
    {

        tmp_node = rb_next(node);
        goal_node = rb_entry(node, user_rb_t, user_node);
        rb_erase(node, root); 
        kfree(goal_node);
        node = tmp_node;
    }
}

struct mytype *my_search(struct rb_root *root, char *string)
{
  struct rb_node *node = root->rb_node;
  while (node)
  {
    struct mytype *data = container_of(node, struct mytype, node);
    int result = strcmp(string, data->keystring);
    if (result < 0)
      node = node->rb_left;
    else if (result > 0)
      node = node->rb_right;
    else
      return data;
  }
  return NULL;
}

插入节点 ，需要先查找到适合插入的位置，然后插入节点，最后调整树的平衡状态

int my_insert(struct rb_root *root, struct mytype *data)
{
  struct rb_node **new = &(root->rb_node), *parent = NULL;

  /* Figure out where to put new node */
  while (*new)
  {
    struct mytype *this = container_of(*new, struct mytype, node);
    int result = strcmp(data->keystring, this->keystring);

    parent = *new;
    if (result < 0)
      new = &((*new)->rb_left);
    else if (result > 0)
      new = &((*new)->rb_right);
    else
      return FALSE;
  }

  /* Add new node and rebalance tree. */
  rb_link_node(&data->node, parent, new);//插入节点时需要把新节点指向其父亲节点,通过该函数完成
  rb_insert_color(&data->node, root);//对新插入的节点着色，并修正红黑树使其达到平衡

  return TRUE;
}

删除节点，直接调用接口

struct mytype *data = mysearch(&mytree, "walrus");
if (data)
{
  rb_erase(&data->node, &mytree); //删除节点，该函数内部调用____rb_erase_color()函数来对红黑树进行修正
  myfree(data);
}

遍历红黑树

struct rb_node *node;
for (node = rb_first(&mytree); node; node = rb_next(node))
  printk("key=%s\n", rb_entry(node, struct mytype, node)->keystring);