blob: fd09465d82ca61e0e73c8b8eb10f085f372f0d1f [file] [log] [blame]
Michel Lespinasse910a7422012-10-08 16:30:39 -07001#include <linux/module.h>
2#include <linux/rbtree.h>
3#include <linux/random.h>
4#include <asm/timex.h>
5
6#define NODES 100
7#define PERF_LOOPS 100000
8#define CHECK_LOOPS 100
9
10struct test_node {
11 struct rb_node rb;
12 u32 key;
13};
14
15static struct rb_root root = RB_ROOT;
16static struct test_node nodes[NODES];
17
18static struct rnd_state rnd;
19
20static void insert(struct test_node *node, struct rb_root *root)
21{
22 struct rb_node **new = &root->rb_node, *parent = NULL;
23
24 while (*new) {
25 parent = *new;
26 if (node->key < rb_entry(parent, struct test_node, rb)->key)
27 new = &parent->rb_left;
28 else
29 new = &parent->rb_right;
30 }
31
32 rb_link_node(&node->rb, parent, new);
33 rb_insert_color(&node->rb, root);
34}
35
36static inline void erase(struct test_node *node, struct rb_root *root)
37{
38 rb_erase(&node->rb, root);
39}
40
41static void init(void)
42{
43 int i;
44 for (i = 0; i < NODES; i++)
45 nodes[i].key = prandom32(&rnd);
46}
47
48static bool is_red(struct rb_node *rb)
49{
50 return !(rb->__rb_parent_color & 1);
51}
52
53static int black_path_count(struct rb_node *rb)
54{
55 int count;
56 for (count = 0; rb; rb = rb_parent(rb))
57 count += !is_red(rb);
58 return count;
59}
60
61static void check(int nr_nodes)
62{
63 struct rb_node *rb;
64 int count = 0;
65 int blacks;
66 u32 prev_key = 0;
67
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))));
73 if (!count)
74 blacks = black_path_count(rb);
75 else
76 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
77 blacks != black_path_count(rb));
78 prev_key = node->key;
79 count++;
80 }
81 WARN_ON_ONCE(count != nr_nodes);
82}
83
84static int rbtree_test_init(void)
85{
86 int i, j;
87 cycles_t time1, time2, time;
88
89 printk(KERN_ALERT "rbtree testing");
90
Michel Lespinasse28d75302012-10-08 16:31:04 -070091 prandom32_seed(&rnd, 3141592653589793238ULL);
Michel Lespinasse910a7422012-10-08 16:30:39 -070092 init();
93
94 time1 = get_cycles();
95
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);
101 }
102
103 time2 = get_cycles();
104 time = time2 - time1;
105
106 time = div_u64(time, PERF_LOOPS);
107 printk(" -> %llu cycles\n", (unsigned long long)time);
108
109 for (i = 0; i < CHECK_LOOPS; i++) {
110 init();
111 for (j = 0; j < NODES; j++) {
112 check(j);
113 insert(nodes + j, &root);
114 }
115 for (j = 0; j < NODES; j++) {
116 check(NODES - j);
117 erase(nodes + j, &root);
118 }
119 check(0);
120 }
121
122 return -EAGAIN; /* Fail will directly unload the module */
123}
124
125static void rbtree_test_exit(void)
126{
127 printk(KERN_ALERT "test exit\n");
128}
129
130module_init(rbtree_test_init)
131module_exit(rbtree_test_exit)
132
133MODULE_LICENSE("GPL");
134MODULE_AUTHOR("Michel Lespinasse");
135MODULE_DESCRIPTION("Red Black Tree test");