Linux内核中红黑树的使用方法

  Linux内核中红黑树的使用方法可以参考：linux/Documentation/rbtree.txt文件。本文的Linux内核以4.1.15版本为例。
  红黑树是一种自平衡二叉搜索树，用于存储可排序的键/值数据对。 这与根树不同是用来有效地存储稀疏数组，从而使用长整数索引插入/访问/删除节点)和哈希表(它们没有被排序容易按顺序遍历，并且必须针对特定的大小和 哈希函数，rbtrees可以优雅地扩展存储任意键)。

一、调用Linux内核中红黑树API时常用的操作

1.1、创建一个新的红黑树

  红黑树中的数据节点是包含结构体rb_node成员的结构 ：

  struct mytype {
  	struct rb_node node;
  	char *keystring;
  };

  当处理指向内嵌结构体rb_node的指针时，包含的数据 结构可以通过标准的container_of()宏来访问。 此外, 单个成员可以通过rb_entry(node, type, member)直接访问。
  在每个红黑树的根上都有一个rb_root结构，初始化为空：

struct rb_root mytree = RB_ROOT;

1.2、在红黑树中搜索值

  为红黑树编写搜索函数相当简单:从根处开始，比较每个值，必要时跟随左分支或右分支。

  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;

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

1.3、在红黑树中插入数据

  在红黑树中插入数据，首先需要搜索新节点要插入的位置，然后插入节点并重新平衡(“重新着色”)树。
  插入的搜索与前面的搜索不同，它查找要在其上移植新节点的指针的位置。 新节点也是需要一个到其父节点的链接以实现重新平衡。

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

1.4、在红黑树中删除数据

  在红黑树中删除数据，需要调用函数：

void rb_erase(struct rb_node *victim, struct rb_root *tree);

例如：

  struct mytype *data = mysearch(&mytree, "walrus");

  if (data) {
  	rb_erase(&data->node, &mytree);
  	myfree(data);
  }

二、使用实例

2.1、拷贝源文件

$ cp ../linux/lib/rbtree.c .
$ cp ../linux/include/linux/rbtree.h

2.2、修改源文件

2.2.1、修改rbtree.c文件

  修改包含头文件的代码

//删除以下两行代码
#include <linux/rbtree.h>
#include <linux/module.h>
//新增以下代码，即包含当前目录中的头文件rbtree.h
#include "rbtree.h"

  删除所有的EXPORT_SYMBOL宏

EXPORT_SYMBOL(rb_insert_color);
EXPORT_SYMBOL(rb_erase);
EXPORT_SYMBOL(rb_augment_insert);
EXPORT_SYMBOL(rb_augment_erase_begin);
EXPORT_SYMBOL(rb_augment_erase_end);
EXPORT_SYMBOL(rb_first);
EXPORT_SYMBOL(rb_last);
EXPORT_SYMBOL(rb_next);
EXPORT_SYMBOL(rb_prev);
EXPORT_SYMBOL(rb_replace_node);

2.2.2、修改rbtree.h文件

//删除以下两行代码
#include <linux/kernel.h>
#include <linux/stddef.h>
 
/* linux/include/linux/stddef.h */
#undef NULL
#if defined(__cplusplus)
#define NULL 0
#else
#define NULL ((void *)0)
#endif
 
/* linux/include/linux/stddef.h */
#define offsetof(TYPE, MEMBER) ((size_t) &((TYPE *)0)->MEMBER)
 
/* linux/include/linux/kernel.h */
#define container_of(ptr, type, member) ({			\
	const typeof( ((type *)0)->member ) *__mptr = (ptr);	\
	(type *)( (char *)__mptr - offsetof(type,member) );})

2.3、示例代码

/* test.c */
#include <stdio.h>
#include <stdlib.h>
#include "rbtree.h"
 
struct mytype {
    struct rb_node my_node;
    int num;
};
 
struct mytype *my_search(struct rb_root *root, int num)
{
    struct rb_node *node = root->rb_node;
 
    while (node) {
	struct mytype *data = container_of(node, struct mytype, my_node);
 
	if (num < data->num)
	    node = node->rb_left;
	else if (num > data->num)
	    node = node->rb_right;
	else
	    return data;
    }
    
    return NULL;
}
 
int my_insert(struct rb_root *root, struct mytype *data)
{
    struct rb_node **tmp = &(root->rb_node), *parent = NULL;
 
    /* Figure out where to put new node */
    while (*tmp) {
	struct mytype *this = container_of(*tmp, struct mytype, my_node);
 
	parent = *tmp;
	if (data->num < this->num)
	    tmp = &((*tmp)->rb_left);
	else if (data->num > this->num)
	    tmp = &((*tmp)->rb_right);
	else 
	    return -1;
    }
    
    /* Add new node and rebalance tree. */
    rb_link_node(&data->my_node, parent, tmp);
    rb_insert_color(&data->my_node, root);
    
    return 0;
}
 
void my_delete(struct rb_root *root, int num)
{
    struct mytype *data = my_search(root, num);
    if (!data) { 
	fprintf(stderr, "Not found %d.\n", num);
	return;
    }
    
    rb_erase(&data->my_node, root);
    free(data);
}
 
void print_rbtree(struct rb_root *tree)
{
    struct rb_node *node;
    
    for (node = rb_first(tree); node; node = rb_next(node))
	printf("%d ", rb_entry(node, struct mytype, my_node)->num);
    
    printf("\n");
}
 
int main(int argc, char *argv[])
{
    struct rb_root mytree = RB_ROOT;
    int i, ret, num;
    struct mytype *tmp;
 
    if (argc < 2) {
	fprintf(stderr, "Usage: %s num\n", argv[0]);
	exit(-1);
    }
 
    num = atoi(argv[1]);
 
    printf("Please enter %d integers:\n", num);
    for (i = 0; i < num; i++) {
	tmp = malloc(sizeof(struct mytype));
	if (!tmp)
	    perror("Allocate dynamic memory");
 
	scanf("%d", &tmp->num);
	
	ret = my_insert(&mytree, tmp);
	if (ret < 0) {
	    fprintf(stderr, "The %d already exists.\n", tmp->num);
	    free(tmp);
	}
    }
 
    printf("\nthe first test\n");
    print_rbtree(&mytree);
 
    my_delete(&mytree, 21);
 
    printf("\nthe second test\n");
    print_rbtree(&mytree);
 
    return 0;
}

2.4、编译并执行

$ gcc rbtree.c test.c -o test
richard@tanglinux:~/algorithm/redblack$ ./test 10
Please enter 10 integers:
23
4
56
32
89
122
12
21
45
23
The 23 already exists.
 
the first test
4 12 21 23 32 45 56 89 122 
 
the second test
4 12 23 32 45 56 89 122