blob: 122f02f9941b5680b6fd9f14793e8806476bb634 [file] [log] [blame]
Michel Lespinasse910a7422012-10-08 16:30:39 -07001#include <linux/module.h>
Michel Lespinasse9c079ad2012-10-08 16:31:33 -07002#include <linux/rbtree_augmented.h>
Michel Lespinasse910a7422012-10-08 16:30:39 -07003#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;
Michel Lespinassedadf9352012-10-08 16:31:15 -070013
14 /* following fields used for testing augmented rbtree functionality */
15 u32 val;
16 u32 augmented;
Michel Lespinasse910a7422012-10-08 16:30:39 -070017};
18
19static struct rb_root root = RB_ROOT;
20static struct test_node nodes[NODES];
21
22static struct rnd_state rnd;
23
24static void insert(struct test_node *node, struct rb_root *root)
25{
26 struct rb_node **new = &root->rb_node, *parent = NULL;
Michel Lespinassedadf9352012-10-08 16:31:15 -070027 u32 key = node->key;
Michel Lespinasse910a7422012-10-08 16:30:39 -070028
29 while (*new) {
30 parent = *new;
Michel Lespinassedadf9352012-10-08 16:31:15 -070031 if (key < rb_entry(parent, struct test_node, rb)->key)
Michel Lespinasse910a7422012-10-08 16:30:39 -070032 new = &parent->rb_left;
33 else
34 new = &parent->rb_right;
35 }
36
37 rb_link_node(&node->rb, parent, new);
38 rb_insert_color(&node->rb, root);
39}
40
41static inline void erase(struct test_node *node, struct rb_root *root)
42{
43 rb_erase(&node->rb, root);
44}
45
Michel Lespinassedadf9352012-10-08 16:31:15 -070046static inline u32 augment_recompute(struct test_node *node)
47{
48 u32 max = node->val, child_augmented;
49 if (node->rb.rb_left) {
50 child_augmented = rb_entry(node->rb.rb_left, struct test_node,
51 rb)->augmented;
52 if (max < child_augmented)
53 max = child_augmented;
54 }
55 if (node->rb.rb_right) {
56 child_augmented = rb_entry(node->rb.rb_right, struct test_node,
57 rb)->augmented;
58 if (max < child_augmented)
59 max = child_augmented;
60 }
61 return max;
62}
63
Michel Lespinasse39088362012-10-08 16:31:21 -070064RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
65 u32, augmented, augment_recompute)
Michel Lespinasse14b94af2012-10-08 16:31:17 -070066
Michel Lespinassedadf9352012-10-08 16:31:15 -070067static void insert_augmented(struct test_node *node, struct rb_root *root)
68{
Michel Lespinasse14b94af2012-10-08 16:31:17 -070069 struct rb_node **new = &root->rb_node, *rb_parent = NULL;
Michel Lespinassedadf9352012-10-08 16:31:15 -070070 u32 key = node->key;
Michel Lespinasse14b94af2012-10-08 16:31:17 -070071 u32 val = node->val;
72 struct test_node *parent;
Michel Lespinassedadf9352012-10-08 16:31:15 -070073
74 while (*new) {
Michel Lespinasse14b94af2012-10-08 16:31:17 -070075 rb_parent = *new;
76 parent = rb_entry(rb_parent, struct test_node, rb);
77 if (parent->augmented < val)
78 parent->augmented = val;
79 if (key < parent->key)
80 new = &parent->rb.rb_left;
Michel Lespinassedadf9352012-10-08 16:31:15 -070081 else
Michel Lespinasse14b94af2012-10-08 16:31:17 -070082 new = &parent->rb.rb_right;
Michel Lespinassedadf9352012-10-08 16:31:15 -070083 }
84
Michel Lespinasse14b94af2012-10-08 16:31:17 -070085 node->augmented = val;
86 rb_link_node(&node->rb, rb_parent, new);
87 rb_insert_augmented(&node->rb, root, &augment_callbacks);
Michel Lespinassedadf9352012-10-08 16:31:15 -070088}
89
90static void erase_augmented(struct test_node *node, struct rb_root *root)
91{
Michel Lespinasse14b94af2012-10-08 16:31:17 -070092 rb_erase_augmented(&node->rb, root, &augment_callbacks);
Michel Lespinassedadf9352012-10-08 16:31:15 -070093}
94
Michel Lespinasse910a7422012-10-08 16:30:39 -070095static void init(void)
96{
97 int i;
Michel Lespinassedadf9352012-10-08 16:31:15 -070098 for (i = 0; i < NODES; i++) {
Akinobu Mita496f2f92012-12-17 16:04:23 -080099 nodes[i].key = prandom_u32_state(&rnd);
100 nodes[i].val = prandom_u32_state(&rnd);
Michel Lespinassedadf9352012-10-08 16:31:15 -0700101 }
Michel Lespinasse910a7422012-10-08 16:30:39 -0700102}
103
104static bool is_red(struct rb_node *rb)
105{
106 return !(rb->__rb_parent_color & 1);
107}
108
109static int black_path_count(struct rb_node *rb)
110{
111 int count;
112 for (count = 0; rb; rb = rb_parent(rb))
113 count += !is_red(rb);
114 return count;
115}
116
117static void check(int nr_nodes)
118{
119 struct rb_node *rb;
Davidlohr Bueso4130f0e2013-04-30 15:28:24 -0700120 int count = 0, blacks = 0;
Michel Lespinasse910a7422012-10-08 16:30:39 -0700121 u32 prev_key = 0;
122
123 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
124 struct test_node *node = rb_entry(rb, struct test_node, rb);
125 WARN_ON_ONCE(node->key < prev_key);
126 WARN_ON_ONCE(is_red(rb) &&
127 (!rb_parent(rb) || is_red(rb_parent(rb))));
128 if (!count)
129 blacks = black_path_count(rb);
130 else
131 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
132 blacks != black_path_count(rb));
133 prev_key = node->key;
134 count++;
135 }
Davidlohr Bueso4130f0e2013-04-30 15:28:24 -0700136
Michel Lespinasse910a7422012-10-08 16:30:39 -0700137 WARN_ON_ONCE(count != nr_nodes);
Davidlohr Bueso4130f0e2013-04-30 15:28:24 -0700138 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);
Michel Lespinasse910a7422012-10-08 16:30:39 -0700139}
140
Michel Lespinassedadf9352012-10-08 16:31:15 -0700141static void check_augmented(int nr_nodes)
142{
143 struct rb_node *rb;
144
145 check(nr_nodes);
146 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
147 struct test_node *node = rb_entry(rb, struct test_node, rb);
148 WARN_ON_ONCE(node->augmented != augment_recompute(node));
149 }
150}
151
Davidlohr Buesoc75aaa82013-04-30 15:28:25 -0700152static int __init rbtree_test_init(void)
Michel Lespinasse910a7422012-10-08 16:30:39 -0700153{
154 int i, j;
155 cycles_t time1, time2, time;
156
157 printk(KERN_ALERT "rbtree testing");
158
Akinobu Mita496f2f92012-12-17 16:04:23 -0800159 prandom_seed_state(&rnd, 3141592653589793238ULL);
Michel Lespinasse910a7422012-10-08 16:30:39 -0700160 init();
161
162 time1 = get_cycles();
163
164 for (i = 0; i < PERF_LOOPS; i++) {
165 for (j = 0; j < NODES; j++)
166 insert(nodes + j, &root);
167 for (j = 0; j < NODES; j++)
168 erase(nodes + j, &root);
169 }
170
171 time2 = get_cycles();
172 time = time2 - time1;
173
174 time = div_u64(time, PERF_LOOPS);
175 printk(" -> %llu cycles\n", (unsigned long long)time);
176
177 for (i = 0; i < CHECK_LOOPS; i++) {
178 init();
179 for (j = 0; j < NODES; j++) {
180 check(j);
181 insert(nodes + j, &root);
182 }
183 for (j = 0; j < NODES; j++) {
184 check(NODES - j);
185 erase(nodes + j, &root);
186 }
187 check(0);
188 }
189
Michel Lespinassedadf9352012-10-08 16:31:15 -0700190 printk(KERN_ALERT "augmented rbtree testing");
191
192 init();
193
194 time1 = get_cycles();
195
196 for (i = 0; i < PERF_LOOPS; i++) {
197 for (j = 0; j < NODES; j++)
198 insert_augmented(nodes + j, &root);
199 for (j = 0; j < NODES; j++)
200 erase_augmented(nodes + j, &root);
201 }
202
203 time2 = get_cycles();
204 time = time2 - time1;
205
206 time = div_u64(time, PERF_LOOPS);
207 printk(" -> %llu cycles\n", (unsigned long long)time);
208
209 for (i = 0; i < CHECK_LOOPS; i++) {
210 init();
211 for (j = 0; j < NODES; j++) {
212 check_augmented(j);
213 insert_augmented(nodes + j, &root);
214 }
215 for (j = 0; j < NODES; j++) {
216 check_augmented(NODES - j);
217 erase_augmented(nodes + j, &root);
218 }
219 check_augmented(0);
220 }
221
Michel Lespinasse910a7422012-10-08 16:30:39 -0700222 return -EAGAIN; /* Fail will directly unload the module */
223}
224
Davidlohr Buesoc75aaa82013-04-30 15:28:25 -0700225static void __exit rbtree_test_exit(void)
Michel Lespinasse910a7422012-10-08 16:30:39 -0700226{
227 printk(KERN_ALERT "test exit\n");
228}
229
230module_init(rbtree_test_init)
231module_exit(rbtree_test_exit)
232
233MODULE_LICENSE("GPL");
234MODULE_AUTHOR("Michel Lespinasse");
235MODULE_DESCRIPTION("Red Black Tree test");