blob: 31dd4ccd3baae800a4459088cca01ca560c3a536 [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
Cody P Schafera791a622013-09-11 14:25:17 -0700117static void check_postorder(int nr_nodes)
118{
119 struct rb_node *rb;
120 int count = 0;
121 for (rb = rb_first_postorder(&root); rb; rb = rb_next_postorder(rb))
122 count++;
123
124 WARN_ON_ONCE(count != nr_nodes);
125}
126
Michel Lespinasse910a7422012-10-08 16:30:39 -0700127static void check(int nr_nodes)
128{
129 struct rb_node *rb;
Davidlohr Bueso4130f0e2013-04-30 15:28:24 -0700130 int count = 0, blacks = 0;
Michel Lespinasse910a7422012-10-08 16:30:39 -0700131 u32 prev_key = 0;
132
133 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
134 struct test_node *node = rb_entry(rb, struct test_node, rb);
135 WARN_ON_ONCE(node->key < prev_key);
136 WARN_ON_ONCE(is_red(rb) &&
137 (!rb_parent(rb) || is_red(rb_parent(rb))));
138 if (!count)
139 blacks = black_path_count(rb);
140 else
141 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
142 blacks != black_path_count(rb));
143 prev_key = node->key;
144 count++;
145 }
Davidlohr Bueso4130f0e2013-04-30 15:28:24 -0700146
Michel Lespinasse910a7422012-10-08 16:30:39 -0700147 WARN_ON_ONCE(count != nr_nodes);
Davidlohr Bueso4130f0e2013-04-30 15:28:24 -0700148 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);
Cody P Schafera791a622013-09-11 14:25:17 -0700149
150 check_postorder(nr_nodes);
Michel Lespinasse910a7422012-10-08 16:30:39 -0700151}
152
Michel Lespinassedadf9352012-10-08 16:31:15 -0700153static void check_augmented(int nr_nodes)
154{
155 struct rb_node *rb;
156
157 check(nr_nodes);
158 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
159 struct test_node *node = rb_entry(rb, struct test_node, rb);
160 WARN_ON_ONCE(node->augmented != augment_recompute(node));
161 }
162}
163
Davidlohr Buesoc75aaa82013-04-30 15:28:25 -0700164static int __init rbtree_test_init(void)
Michel Lespinasse910a7422012-10-08 16:30:39 -0700165{
166 int i, j;
167 cycles_t time1, time2, time;
168
169 printk(KERN_ALERT "rbtree testing");
170
Akinobu Mita496f2f92012-12-17 16:04:23 -0800171 prandom_seed_state(&rnd, 3141592653589793238ULL);
Michel Lespinasse910a7422012-10-08 16:30:39 -0700172 init();
173
174 time1 = get_cycles();
175
176 for (i = 0; i < PERF_LOOPS; i++) {
177 for (j = 0; j < NODES; j++)
178 insert(nodes + j, &root);
179 for (j = 0; j < NODES; j++)
180 erase(nodes + j, &root);
181 }
182
183 time2 = get_cycles();
184 time = time2 - time1;
185
186 time = div_u64(time, PERF_LOOPS);
187 printk(" -> %llu cycles\n", (unsigned long long)time);
188
189 for (i = 0; i < CHECK_LOOPS; i++) {
190 init();
191 for (j = 0; j < NODES; j++) {
192 check(j);
193 insert(nodes + j, &root);
194 }
195 for (j = 0; j < NODES; j++) {
196 check(NODES - j);
197 erase(nodes + j, &root);
198 }
199 check(0);
200 }
201
Michel Lespinassedadf9352012-10-08 16:31:15 -0700202 printk(KERN_ALERT "augmented rbtree testing");
203
204 init();
205
206 time1 = get_cycles();
207
208 for (i = 0; i < PERF_LOOPS; i++) {
209 for (j = 0; j < NODES; j++)
210 insert_augmented(nodes + j, &root);
211 for (j = 0; j < NODES; j++)
212 erase_augmented(nodes + j, &root);
213 }
214
215 time2 = get_cycles();
216 time = time2 - time1;
217
218 time = div_u64(time, PERF_LOOPS);
219 printk(" -> %llu cycles\n", (unsigned long long)time);
220
221 for (i = 0; i < CHECK_LOOPS; i++) {
222 init();
223 for (j = 0; j < NODES; j++) {
224 check_augmented(j);
225 insert_augmented(nodes + j, &root);
226 }
227 for (j = 0; j < NODES; j++) {
228 check_augmented(NODES - j);
229 erase_augmented(nodes + j, &root);
230 }
231 check_augmented(0);
232 }
233
Michel Lespinasse910a7422012-10-08 16:30:39 -0700234 return -EAGAIN; /* Fail will directly unload the module */
235}
236
Davidlohr Buesoc75aaa82013-04-30 15:28:25 -0700237static void __exit rbtree_test_exit(void)
Michel Lespinasse910a7422012-10-08 16:30:39 -0700238{
239 printk(KERN_ALERT "test exit\n");
240}
241
242module_init(rbtree_test_init)
243module_exit(rbtree_test_exit)
244
245MODULE_LICENSE("GPL");
246MODULE_AUTHOR("Michel Lespinasse");
247MODULE_DESCRIPTION("Red Black Tree test");