Linux核心中的紅/黑樹狀結構的使用

最近需要使用紅/黑樹狀結構，在網上尋找資料的時候無意中發現linux核心中有個紅/黑樹狀結構的實現，並且其代碼非常的獨立，現把它摘錄出來。我摘錄自2.6.24的核心，分為兩個檔案rbtree.h和rbtree.c，rbtree.h位於核心源碼的include/linux目錄中，rbtree.c位於核心源碼的lib目錄中。

rbtree.h中刪除#include和#include兩行，添加#include

對於rb_node的聲明刪除掉最後的__attribute__((aligned(sizeof(long))))

最後在rbtree.h中添加如下一些宏定義：

#ifdef __compiler_offsetof
#define offsetof(TYPE,MEMBER) __compiler_offsetof(TYPE,MEMBER)
#else
#define offsetof(TYPE, MEMBER) ((size_t) &((TYPE *)0)->MEMBER)
#endif

#define container_of(ptr, type, member) ({            /
        const typeof( ((type *)0)->member ) *__mptr = (ptr);    /
        (type *)( (char *)__mptr - offsetof(type,member) );})

 

在使用rbtree的時候，每個rbtree的元素都由使用者自己定義如：
struct my_item
{
  struct rb_node node;
  struct my_data_struct data;
};
其中的struct my_data_struct是自己的資料區域
struct my_data_struct
{
  long key;
  long value;
};

下面需要實現兩個函數，第一個是一個search函數，形式如下：
struct my_item* _my_rbtree_search (struct rb_root* root, long key);
其中的root參數是根節點指標；key是關鍵字，rbtree根據它來排序以及尋找。
該函數的實現如下：
struct my_item*
_my_rbtree_search (struct rb_root* root, long key)
{
  struct rb_node *node;
  struct my_item* item;
  node = root->rb_node;

  while (node)
    {
      item = rb_entry (node, struct my_item, node);
      if (item->data.key > key)
        node = node->rb_left;
      else if (item->data.key         node = node->rb_right;
      else
        {
          return item;            /* found it */
        }
    }
  return NULL;
}

下面還需要實現一個插入的函數：
void
_my_rbtree_insert (struct rb_root* root, struct my_item* item)
{
  struct rb_node **link, *parent;
  long value;
  struct my_item* p_item;
  link = &(root->rb_node);
  value = item->data.key;
  /* Go to the bottom of the tree */
  while (*link)
    {
      parent = *link;
      p_item = rb_entry(parent, struct my_item, node);
      if (p_item->data.key > value)
        link = &((*link)->rb_left);
      else if (p_item->data.key         link = &((*link)->rb_right);
      else
        return;
    }
  /* Put the new node there */
  rb_link_node(&(item->node), parent, link);
  rb_insert_color(&(item->node), root);
}

擁有了以上兩個函數我們就可以實現search、insert、remove、modify、traverse等操作，下面給出大概的虛擬碼吧
struct my_data_struct search (struct rb_root* root, long key)
{
  struct my_data_struct data;
  struct my_item *item;
  memset (&data, 0, sizeof(struct my_data));
  item = _my_rbtree_search (root, key);
  if (NULL != item)
  {
    memcpy (&data, item->data, sizeof(struct my_data_struct));
  }
  return data;
}

int insert (struct rb_root* root, long key, long value)
{
  struct my_item  *item = (struct my_item*) malloc (sizeof(struct my_item));
  if (NULL == item)
  {
    return -1;
  }
  item->data.key = key;
  item->data.value = value;
  _my_rbtree_insert (root, item);
  return 0;
}

int remove (struct rb_root* root, long key)
{
  struct my_item* item;
  item = _my_rbtree_search (root, key);
  if (NULL != item)
  {
    rb_erase (&(item->node), root);
  }
  return 0;
}

int modify (struct rb_root* root, long key, long new_value)
{
  struct my_item* item;
  item = _my_rbtree_search (root, key);
  if (NULL != item)
  {
    item->data.value = new_value;
  }
  return 0;
}

int traverse (struct rb_root* root)
{
  struct rb_node* node;
  struct my_item* item;
  for (node = rb_first (root); node; node = rb_next (node))
  {
    item = rb_entry (node, struct my_item, node);
    printf ("key:%d/tvalue:%d/n", item->data.key, item->data.value);
  }
  return 0;
}