rbtree原理及应用--使用

kernel中红黑树数据结构的使用方法

#include <ds/rbtree_augmented.h>
#include <types.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#define MAX_INT		40
#define MAC_TEST 0
#define INT_TEST 1
#define TEST_INSERT		0
#define TEST_DELETE		1

struct user_rb_t {
	struct rb_node rbnode;
	u8 mac[ETH_ALEN];
	int v;
};

struct rb_root root = RB_ROOT;

//  创建节点
struct user_rb_t * create_user_node(u8 *mac, int v)
{
	struct user_rb_t *node;
	
	node = (struct user_rb_t *)malloc(sizeof(struct user_rb_t));
	if (NULL == node)
		return NULL;
	memset(node, 0, sizeof(struct user_rb_t));
#if MAC_TEST
	memcpy(node->mac, mac, ETH_ALEN);
#elif INT_TEST
	node->v = v;
#endif
	return node;
}

// 查找节点
struct user_rb_t * user_node_find(struct rb_root *root, u8 *mac, int v)
{
	struct rb_node *node = root->rb_node;
	struct user_rb_t * target = NULL;
	int res = 0;

	while (node) {
		target = rb_entry(node, struct user_rb_t, rbnode);
		#if MAC_TEST
		res = memcmp(target->mac, mac, ETH_ALEN);
		#elif INT_TEST
		res = target->v - v;
		#endif
		if (res < 0) {
			node = node->rb_left;
		} else if (res > 0) {
			node = node->rb_right;
		} else
			return target;
	}
	return NULL;
}

// 插入节点
int user_rb_insert(struct rb_root *root, u8 *mac, int v)
{
	struct rb_node **new = &(root->rb_node);
	struct rb_node *parent = NULL;
	struct user_rb_t *target = NULL;
	int res = 0;

	while (*new) {
		parent = *new;
		target = rb_entry(*new, struct user_rb_t, rbnode);
		#if MAC_TEST
		res = memcmp(target->mac, mac, ETH_ALEN);
		#elif INT_TEST
		res = target->v - v;
		#endif
		if (res < 0) {
			new = &((*new)->rb_left);
		} else if (res > 0) {
			new = &((*new)->rb_right);
		} else
			return true;
	}
	target = create_user_node(mac, v);
	if (NULL == target)
		return false;
	// add new node
	rb_link_node(&target->rbnode, parent, new);
	// rebalance rbtree
	rb_insert_color(&target->rbnode, root);
	return true;
}

// 删除单个节点
void user_rb_delete(struct rb_root *root, u8 *mac, int v)
{
	struct user_rb_t *target = user_node_find(root, mac, v);
	if (NULL != target) {
		rb_erase(&target->rbnode, root);
		free(target);
	}
}

// 删除整棵红黑树
void user_rb_destroy(struct rb_root *root)
{
	struct rb_node *node = NULL, *tmp_node = NULL;
	struct user_rb_t * target = NULL;

	for (node = rb_first(root); node;) {
		tmp_node = rb_next(node);
		target = rb_entry(node, struct user_rb_t, rbnode);
		rb_erase(node, root);
		free(target);
		node = tmp_node;
	}
}

// 遍历红黑树
static void tranverse(struct rb_root *root)
{
	struct rb_node *node;
	int ind = 0;
	for (node = rb_first(root); node; node = rb_next(node)) {
		struct user_rb_t *target = rb_entry(node, struct user_rb_t, rbnode);
		#if MAC_TEST
		printf("\t[%d] key: %02x:%02x:%02x:%02x:%02x:%02x\n", ind,
				target->mac[0], target->mac[1], target->mac[2], 
				target->mac[3], target->mac[4], target->mac[5]);
		#elif INT_TEST
		printf("\t[%d] %p key: %d, pc: 0x%x [ %c ], rb_left: %p, rb_right: %p\n",
				ind, target, target->v,
				target->rbnode.__rb_parent_color,
				rb_is_black(&target->rbnode) ? 'B' : 'R',
				target->rbnode.rb_left,
				target->rbnode.rb_right);
		#endif
		ind++;
		usleep(100000);
	}
}

int gen_imac(int * imac)
{
	unsigned char v = 0;
	int fd = 0;
	char buf[3] = {0};
	int re = 0;
	unsigned short vv = 0;

	#define RANDOM_DEV		"/dev/random"
	fd = open(RANDOM_DEV, O_RDONLY);
	if (-1 != fd) {
		int i = 0;
		for (i = 0; i < MAX_INT; i++) {
			re = read(fd, &vv, 2);
			imac[i] = vv;
		}
		close(fd);
	}
}

#define rebuild_rbtree()	do { \
	user_rb_destroy(&root); \
	int i = 0; \
	for (i = 0; i < MAX_INT; i++) { \
		user_rb_insert(&root, mac[i], imac[i]); \
	} \
} while (0)

int main(int argc, char ** argv)
{
	int ret = 0;
	int imac_bk[MAX_INT] = {0};
	u8 mac[11][ETH_ALEN] = {
		{ 0xea, 0x2a, 0xea, 0xc1, 0xf3, 0x36 },
		{ 0x00, 0x26, 0xc7, 0x86, 0xa5, 0x0e },
		{ 0x1c, 0x65, 0x9d, 0xd9, 0xed, 0x09 },
		{ 0x20, 0xa6, 0x80, 0x15, 0xf4, 0x76 },
		{ 0x20, 0xa6, 0x80, 0x15, 0xf4, 0x76 },
		{ 0x64, 0xa6, 0x51, 0x56, 0x4b, 0xe5 },
		{ 0xb0, 0xe2, 0x35, 0x22, 0x09, 0x50 },
		{ 0x5c, 0xf3, 0xfc, 0x49, 0xae, 0x74 },
		{ 0xC6, 0x93, 0x44, 0x01, 0x01, 0x09 },
		{ 0x5e, 0xf3, 0xfc, 0x4f, 0xae, 0x77 },
		{ 0xB8, 0x2A, 0x72, 0xDA, 0x3E, 0x57 }
	};
	int imac[10] = { 0xea, 0x00, 0x1c, 0x20, 0x2a, 0x64, 0xb0, /*0x5c,*/130, 0xc6, /*0x5e,*/ /*0xb8,*/ 200 };
				  /*  234,    0,   28,   32,   42,  100,  176,   92,  198, 	 94,  184 */
	int i = 0; 
	int first = 1;


	//gen_imac(imac);
	imac[0] = 10;
	imac[1] = 6;
	imac[2] = 1;
	imac[3] = 3;
	imac[4] = 7;
	imac[5] = 2;
	imac[6] = 5;
	imac[7] = 11;
	#if TEST_INSERT
	for (i = 0; i < 8; i++) {
		printf("------     BEGIN %d     ------\n", i);
		user_rb_insert(&root, mac[i], imac[i]);
		tranverse(&root);
		printf("------      END     ------\n");
	}
	#endif
	#if TEST_DELETE
	imac[0] = 1;
	imac[1] = 33;
	imac[2] = 15;
	imac[3] = 26;
	imac[4] = 12;
	imac[5] = 38;
	imac[6] = 28;
	imac[7] = 6;
	imac[8] = 8;
	imac[9] = 36;
	imac[10] = 18;
	imac[11] = 20;
	imac[12] = 4;
	imac[13] = 17;
	imac[14] = 9;
	imac[15] = 14;
	imac[16] = 37;
	imac[17] = 19;
	imac[18] = 10;
	imac[19] = 25;
	imac[20] = 40;
	imac[21] = 3;
	imac[22] = 11;
	imac[23] = 24;
	imac[24] = 30;
	imac[25] = 35;
	imac[26] = 22;
	imac[27] = 5;
	imac[28] = 31;
	imac[29] = 23;
	imac[30] = 32;
	imac[31] = 21;
	imac[32] = 1;
	imac[33] = 33;
	imac[34] = 15;
	imac[35] = 26;
	imac[36] = 12;
	imac[37] = 39;
	imac[38] = 28;
	imac[39] = 6;
	memcpy(imac_bk, imac, sizeof(imac));
	for (i = 0; i < MAX_INT; i++) {
		imac[i] *= 4;
	}
	for (i = 0; i < MAX_INT; i++) {
		del_print_f = 0;
		rebuild_rbtree();
		user_rb_insert(&root, mac[i], 15);
		del_print_f = 1;
		if (1 == first) {
			tranverse(&root);
			first = 0;
		}
		printf("------   DELETE %d/%d   ------\n", imac[i], i);
		user_rb_delete(&root, NULL, imac[i]);
		printf("------   END DEL %d/%d   ------\n", imac[i], i);
	}
	#endif
	
	return 0;
}