1 #include <linux/module.h>
2 #include <linux/rbtree.h>
3 #include <linux/random.h>
7 #define PERF_LOOPS 100000
8 #define CHECK_LOOPS 100
15 static struct rb_root root = RB_ROOT;
16 static struct test_node nodes[NODES];
18 static struct rnd_state rnd;
20 static void insert(struct test_node *node, struct rb_root *root)
22 struct rb_node **new = &root->rb_node, *parent = NULL;
26 if (node->key < rb_entry(parent, struct test_node, rb)->key)
27 new = &parent->rb_left;
29 new = &parent->rb_right;
32 rb_link_node(&node->rb, parent, new);
33 rb_insert_color(&node->rb, root);
36 static inline void erase(struct test_node *node, struct rb_root *root)
38 rb_erase(&node->rb, root);
41 static void init(void)
44 for (i = 0; i < NODES; i++)
45 nodes[i].key = prandom32(&rnd);
48 static bool is_red(struct rb_node *rb)
50 return !(rb->__rb_parent_color & 1);
53 static int black_path_count(struct rb_node *rb)
56 for (count = 0; rb; rb = rb_parent(rb))
61 static void check(int nr_nodes)
68 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
69 struct test_node *node = rb_entry(rb, struct test_node, rb);
70 WARN_ON_ONCE(node->key < prev_key);
71 WARN_ON_ONCE(is_red(rb) &&
72 (!rb_parent(rb) || is_red(rb_parent(rb))));
74 blacks = black_path_count(rb);
76 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
77 blacks != black_path_count(rb));
81 WARN_ON_ONCE(count != nr_nodes);
84 static int rbtree_test_init(void)
87 cycles_t time1, time2, time;
89 printk(KERN_ALERT "rbtree testing");
91 prandom32_seed(&rnd, 3141592653589793238ULL);
96 for (i = 0; i < PERF_LOOPS; i++) {
97 for (j = 0; j < NODES; j++)
98 insert(nodes + j, &root);
99 for (j = 0; j < NODES; j++)
100 erase(nodes + j, &root);
103 time2 = get_cycles();
104 time = time2 - time1;
106 time = div_u64(time, PERF_LOOPS);
107 printk(" -> %llu cycles\n", (unsigned long long)time);
109 for (i = 0; i < CHECK_LOOPS; i++) {
111 for (j = 0; j < NODES; j++) {
113 insert(nodes + j, &root);
115 for (j = 0; j < NODES; j++) {
117 erase(nodes + j, &root);
122 return -EAGAIN; /* Fail will directly unload the module */
125 static void rbtree_test_exit(void)
127 printk(KERN_ALERT "test exit\n");
130 module_init(rbtree_test_init)
131 module_exit(rbtree_test_exit)
133 MODULE_LICENSE("GPL");
134 MODULE_AUTHOR("Michel Lespinasse");
135 MODULE_DESCRIPTION("Red Black Tree test");