linux内核中的红黑树

和内核中的hash table一样，内核中的红黑树比较“裸”。出于效率方面的考虑，并没有把所有的操作都封装起来。

在执行插入操作的时候，需要使用者自己按照搜索二叉树的方法找到想要插入的节点的父节点，然后调用rb_link_node函数将节点插入，再调用rb_insert_color函数对红黑树进行适当的“旋转”。

而在搜索红黑树的时候，则需要使用者自己按照普通二叉树的搜索方法进行搜索。当然，如何比较键值的大小也是使用者自己决定的。

内核的Documentation目录下有一篇rbtree.txt的文档详细的介绍了该如何使用红黑树。

下面是我写的一个使用红黑树的小例子，可在2.6.32下运行。在这个小程序中，coef被当作键值来使用。

#include <linux/init.h>
#include <linux/module.h>
#include <linux/rbtree.h>

struct q_coef
{
    u8 coef;
    u8 index;
    struct rb_node node;
};

#define COEF_NUM 15
u8 coef[15] = {
    0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
    0x1d, 0x3a, 0x74, 0xe8, 0xcd, 0x87, 0x13,
};
struct q_coef q_coef[COEF_NUM];

static void q_coef_init(void)
{
    int i;
    memset(&q_coef, 0, sizeof(q_coef));
    for (i = 0 ; i < COEF_NUM ; i++) {
        q_coef[i].coef = coef[i];
        q_coef[i].index = i + 1;
    }
}

struct rb_root q_coef_tree = RB_ROOT;

static int q_coef_insert(struct rb_root *root, struct q_coef *data)
{
    struct rb_node **new = &(root->rb_node), *parent = NULL;

    /* Figure out where to put new code */
    while (*new) {
        struct q_coef *this = rb_entry(*new, struct q_coef, node);
        parent = *new;
        if (data->coef < this->coef)
            new = &((*new)->rb_left);
        else if (data->coef > this->coef)
            new = &((*new)->rb_right);
        else
            return -1;
    }

    /* Add new node and rebalance tree. */
    rb_link_node(&data->node, parent, new);
    rb_insert_color(&data->node, root);

    return 0;
}

static struct q_coef *q_coef_search(struct rb_root *root, u8 coef)
{
    struct rb_node *node = root->rb_node;
    while (node) {
        struct q_coef *data = rb_entry(node, struct q_coef, node);
        if (coef < data->coef)
            node = node->rb_left;
        else if (coef > data->coef)
            node = node->rb_right;
        else
            return data;
    }
    return NULL;
}

static int rbtest_init (void)
{
    int i;
    struct q_coef *ptr;
    struct rb_node *node;
    int ret;

    q_coef_init();

    for (i = 0 ; i < COEF_NUM ; i++) {
        ret = q_coef_insert(&q_coef_tree, &q_coef[i]);
        if (ret < 0) {
            printk(KERN_WARNING "q_coef_insert failed, i=%d\n", i);
            return -1;
        }
    }

    printk(KERN_INFO "search by input order:\n");
    for (i = 0 ; i < COEF_NUM ; i++) {
        ptr = q_coef_search(&q_coef_tree, coef[i]);
        if (ptr == NULL) {
            printk(KERN_WARNING "q_coef_search failed, i=%d\n", i);
            return -1;
        }
        printk(KERN_INFO "coef[%02d]=0x%02x  ptr->coef=0x%02x ptr->index=%02d\n",
            i, coef[i], ptr->coef, ptr->index);
    }

    printk(KERN_INFO "search from first:\n");
    for (node = rb_first(&q_coef_tree) ; node ; node = rb_next(node)) {
        ptr = rb_entry(node, struct q_coef, node);
        printk(KERN_INFO "ptr->coef=0x%02x  ptr->index=%02d\n", ptr->coef, ptr->index);
    }

    printk(KERN_INFO "search from last:\n");
    for (node = rb_last(&q_coef_tree) ; node ; node = rb_prev(node)) {
        ptr = rb_entry(node, struct q_coef, node);
        printk(KERN_INFO "ptr->coef=0x%02x  ptr->index=%02d\n", ptr->coef, ptr->index);
    }

    printk(KERN_INFO "rbtest done\n");
    return -1;
}

static void rbtest_exit (void)
{
}

module_init(rbtest_init);
module_exit(rbtest_exit);

MODULE_LICENSE("Dual BSD/GPL");