Linux内核中rbtree的使用

在结构体内嵌struct rb_node

定义根节点struct rb_root

把节点依据key加入红黑树

移除节点

遍历红黑树

#include <linux/module.h>
#include <linux/rbtree.h>
#include <linux/random.h>

#define NODES       20

struct test_node {
	u32 key;
	struct rb_node rb;

	u32 val;
};

static struct rb_root root = RB_ROOT;
static struct test_node nodes[NODES];

static void insert(struct test_node *node, struct rb_root *root)
{
	struct rb_node **new = &root->rb_node, *parent = NULL;
	u32 key = node->key;

	while (*new) {
		parent = *new;
		if (key < rb_entry(parent, struct test_node, rb)->key)
			 new = &parent->rb_left;
		else
			new = &parent->rb_right;
	}

	rb_link_node(&node->rb, parent, new);
	rb_insert_color(&node->rb, root);
}

static inline void erase(struct test_node *node, struct rb_root *root)
{
	rb_erase(&node->rb, root);
}

static void traverse(struct rb_root *root)
{
	int i = 0;
	struct test_node *tn;
	struct rb_node *node = rb_first(root);

	printk("%s, %d\n", __func__, __LINE__);

	while (node) {
		tn = rb_entry(node, struct test_node, rb);
		printk("[%2d] tn->key: %u,	tn->val: %u\n", i, tn->key,
		       tn->val);
		i++;
		node = rb_next(node);
	}
}

static struct test_node *find(u32 key, struct rb_root *root)
{
	int found = 0;
	struct test_node *tn;
	struct rb_node *node = root->rb_node;

	while (node) {
		tn = rb_entry(node, struct test_node, rb);
		if (key < tn->key) {
			node = node->rb_left;
		} else if (key > tn->key) {
			node = node->rb_right;
		} else {
			found = 1;
			break;
		}
	}

	if (found) {
		return rb_entry(node, struct test_node, rb);
	}

	return NULL;
}

static void init(void)
{
	int i;
	for (i = 0; i < NODES; i++) {
		nodes[i].key = random32() % 1024;
		nodes[i].val = random32() % 1024;
	}
}

static void dump_nodes(void)
{
	int i;
	for (i = 0; i < NODES; i++) {
		printk("nodes[%2d].key: %4u,	nodes[%2d].val: %4u\n",
		       i, nodes[i].key, i, nodes[i].val);
	}
}

static int __init test_init(void)
{
	int i;
	struct test_node *tn;

	init();
	dump_nodes();

	for (i = 0; i < NODES; i++) {
		insert(&nodes[i], &root);
	}
	traverse(&root);

	erase(&nodes[1], &root);
	traverse(&root);

	nodes[1].key = 123;
	insert(&nodes[1], &root);
	traverse(&root);

	tn = find(123, &root);
	if (tn) {
		printk("found, tn->key: %u,	tn->val: %u\n", tn->key, tn->val);
		erase(tn, &root);
		traverse(&root);
	} else {
		printk("key: 123, not found\n");
	}

	tn = find(222, &root);
	if (tn) {
		printk("found, tn->key: %u,	tn->val: %u\n", tn->key, tn->val);
	} else {
		printk("key: 222, not found\n");
	}

	return -EAGAIN;
}

static void __exit test_exit(void)
{

}

module_init(test_init);
module_exit(test_exit);
MODULE_AUTHOR("tony");
MODULE_DESCRIPTION("test");
MODULE_LICENSE("GPL");
MODULE_ALIAS("test");