如何將 Linux 內核實現的紅黑樹 rbtree 運用到你的 C 程序中？

相信大家都知道紅黑樹是什么吧，但是呢......如果你確實不知道，你不該穿越到這兒的，你應該去這里，這里，還有這里看看，然后再來這里看看，最后如果大爺您賞臉，再來看看我吧 :-)

廢話少說，直接入正題吧，Linux 內核為我們實現了簡潔高效但是......卻不那么容易使用的紅黑樹，如何在你的 C 程序里面使用內核開發者為我們實現的紅黑樹呢，別急別急，本文將一一為您呈現。

Linux 內核紅黑樹的實現代碼位於：lib/rbtree.c，同時頭文件在 include/linux/rbtree.h 中，內核中很多模塊都使用了紅黑樹，詳細介紹參見內核文檔 Documentation/rbtree.txt。

內核中紅黑樹定義如下：

struct rb_node
{
    unsigned long  rb_parent_color;
#define    RB_RED        0
#define    RB_BLACK    1
    struct rb_node *rb_right;
    struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));

struct rb_root
{
    struct rb_node *rb_node;
};

在使用內核的紅黑樹時，需將 struct rb_node 結構包含在自己的數據結構中，比如：

  struct mynode {
      struct rb_node node;
      char *string;
      /*more stuff of your structure hereby*/
  };

可以通過container_of宏獲取包含了 rb_node 結構的起始地址，也可以通過rb_entry(node, type, member),其實：

 

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

 

container_of 中使用了offsetof宏，它們都是內核中常見的宏，其定義如下：

#if defined(container_of)
  #undef container_of
  #define container_of(ptr, type, member) ({            \
        const typeof( ((type *)0)->member ) *__mptr = (ptr);    \
        (type *)( (char *)__mptr - offsetof(type,member) );})
#endif

#if defined(offsetof)
  #undef offsetof
  #define offsetof(TYPE, MEMBER) ((size_t) &((TYPE *)0)->MEMBER)
#endif

如果對上面兩個宏的意思不理解，可以在本文后面留言 :-)

為了使用內核提供的紅黑樹，你需要自己實現插入和查找函數，如下：

struct mynode *my_search(struct rb_root *root, char *string)
{
      struct rb_node *node = root->rb_node;

      while (node) {
          struct mynode *data = container_of(node, struct mynode, node);
        int result;

        result = strcmp(string, data->string);

        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 mynode *data)
{
      struct rb_node **new = &(root->rb_node), *parent = NULL;

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

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

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

    return 1;
}

釋放某一節點空間：

void my_free(struct mynode *node)
{
    if (node != NULL) {
        if (node->string != NULL) {
            free(node->string);
            node->string = NULL;
        }
        free(node);
        node = NULL;
    }
}

綜合上面的代碼：

#define NUM_NODES 32

int main()
{

    struct mynode *mn[NUM_NODES];

    /* *insert */
    int i = 0;
    for (; i < NUM_NODES; i++) {
        mn[i] = (struct mynode *)malloc(sizeof(struct mynode));
        mn[i]->string = (char *)malloc(sizeof(char) * 4);
        sprintf(mn[i]->string, "%d", i);
        my_insert(&mytree, mn[i]);
    }
    
    /* *search */
    struct rb_node *node;
    for (node = rb_first(&mytree); node; node = rb_next(node))
        printf("key = %s\n", rb_entry(node, struct mynode, node)->string);

    /* *delete */
    printf("delete node 20: \n");
    struct mynode *data = my_search(&mytree, "20");
    if (data) {
        rb_erase(&data->node, &mytree);
        my_free(data);
    }


    /* *delete again*/
    printf("delete node 10: \n");
    data = my_search(&mytree, "10");
    if (data) {
        rb_erase(&data->node, &mytree);
        my_free(data);
    }


    /* *delete once again*/
    printf("delete node 15: \n");
    data = my_search(&mytree, "15");
    if (data) {
        rb_erase(&data->node, &mytree);
        my_free(data);
    }

    /* *search again*/
    printf("search again:\n");
    for (node = rb_first(&mytree); node; node = rb_next(node))
        printf("key = %s\n", rb_entry(node, struct mynode, node)->string);
    return 0;
}

完整代碼請移步我的github：https://github.com/forhappy/rbtree