linux 内核中的红黑树数据结构举例

/*********************************************************
*
*  linux 内核中的红黑树数据结构举例
*
*      Auth: Alanx
*      Email: [e-mail]
*      2009-10-12
*
*      本例适当修改可直接使用内核中的红黑树
*      include/linux/rbtree.h
*      lib/rbtree.c
*
*      在rbtree.h中增加宏:
*
*      #define offsetof(TYPE, MEMBER) ((size_t) &((TYPE *)0)->MEMBER) 
*      #define container_of(ptr, type, member) ( { \
*            const typeof( ((type *)0)->member ) *__mptr = (ptr); \
*           (type *)( (char *)__mptr - offsetof(type,member) ); } )
*
*      在rbtree.c中去全部 EXPORT_SYMBOL 符中,
*          并适当增减 .h文件
*
*********************************************************/  

#include "rbtree.h"

struct ip_tree
{
unsigned long ip;
struct rb_node tree;
};

struct ip_tree *search_ip(struct rb_root *root, unsigned long ip)
{
struct rb_node *p;
struct ip_tree *ip_node;
p = root->rb_node;
while(p)
{
ip_node = rb_entry(p, struct ip_tree, tree);
if (ip_node->ip > ip)
p = p->rb_left;
else if (ip_node->ip < ip)
p = p->rb_right;
else
return ip_node;
}
return NULL;
}

int init_rb_tree(struct rb_root *root, unsigned long ip)
{

struct ip_tree *ip_node;
ip_node = calloc(sizeof(struct ip_tree), 1);
if (!ip_node)
return 0;
ip_node->ip = ip;
rb_link_node(&(ip_node->tree), NULL, &(root->rb_node));
rb_insert_color(&(ip_node->tree), root);
}

void _insert_node(struct rb_root *root, struct ip_tree  *new_ip_node)
{
struct rb_node **link;
struct rb_node *parent;
struct ip_tree *ip_node;

link = &(root->rb_node);

while(*link)
{
parent = *link;
ip_node = rb_entry(parent, struct ip_tree, tree); 
if(ip_node->ip > new_ip_node->ip)
link = &((*link)->rb_left);
else if(ip_node->ip < new_ip_node->ip)
link = &((*link)->rb_right);
else 
return ;
}
rb_link_node(&(new_ip_node->tree), parent, link);
rb_insert_color(&(new_ip_node->tree), root);
}

int insert_ip(struct rb_root *root, unsigned long ip)
{
struct ip_tree *ip_node;
ip_node = calloc(sizeof(struct ip_tree), 1);
if (!ip_node)
return 0;
ip_node->ip = ip;
_insert_node(root,ip_node);
return 1;
}

int modify_ip(struct rb_root *root, unsigned long ip, unsigned long new_ip)
{
struct ip_tree *ip_node;
ip_node = search_ip(root, ip);

if (ip_node)
{
ip_node->ip = new_ip;
return 1;
}
return 0;
}

void  remove_ip(struct rb_root *root, unsigned long ip)
{
struct ip_tree *ip_node;
ip_node = search_ip(root, ip);

if (ip_node)
{
rb_erase(&(ip_node->tree), root);
free(ip_node);
}
return;
}

void traverse_tree(struct rb_root* root) 
{ 
struct rb_node* p; 
struct ip_tree *ip_node;
unsigned long ip;
for (p = rb_first (root); p; p = rb_next (p)) 
{ 
ip_node = rb_entry (p, struct ip_tree, tree); 
ip = ip_node->ip;
printf ("ip:%d\n",ip);  
} 
}



int main()
{
struct rb_root root = {NULL};
init_rb_tree(&root, 2222);
insert_ip(&root, 1311);
insert_ip(&root, 1234);
insert_ip(&root, 1123);
insert_ip(&root, 1331);
traverse_tree(&root);
printf("------------------------------\n");
modify_ip(&root, 1234, 4321);
traverse_tree(&root);
printf("------------------------------\n");
remove_ip(&root, 1123);
insert_ip(&root, 1124);
traverse_tree(&root);
return 0;
}