1// SPDX-License-Identifier: GPL-2.0-only
2#include <linux/module.h>
3#include <linux/moduleparam.h>
4#include <linux/rbtree_augmented.h>
5#include <linux/prandom.h>
6#include <linux/slab.h>
7#include <asm/timex.h>
8
9#define __param(type, name, init, msg) \
10 static type name = init; \
11 module_param(name, type, 0444); \
12 MODULE_PARM_DESC(name, msg);
13
14__param(int, nnodes, 100, "Number of nodes in the rb-tree");
15__param(int, perf_loops, 1000, "Number of iterations modifying the rb-tree");
16__param(int, check_loops, 100, "Number of iterations modifying and verifying the rb-tree");
17__param(ullong, seed, 3141592653589793238ULL, "Random seed");
18
19struct test_node {
20 u32 key;
21 struct rb_node rb;
22
23 /* following fields used for testing augmented rbtree functionality */
24 u32 val;
25 u32 augmented;
26};
27
28static struct rb_root_cached root = RB_ROOT_CACHED;
29static struct test_node *nodes = NULL;
30
31static struct rnd_state rnd;
32
33static void insert(struct test_node *node, struct rb_root_cached *root)
34{
35 struct rb_node **new = &root->rb_root.rb_node, *parent = NULL;
36 u32 key = node->key;
37
38 while (*new) {
39 parent = *new;
40 if (key < rb_entry(parent, struct test_node, rb)->key)
41 new = &parent->rb_left;
42 else
43 new = &parent->rb_right;
44 }
45
46 rb_link_node(node: &node->rb, parent, rb_link: new);
47 rb_insert_color(&node->rb, &root->rb_root);
48}
49
50static void insert_cached(struct test_node *node, struct rb_root_cached *root)
51{
52 struct rb_node **new = &root->rb_root.rb_node, *parent = NULL;
53 u32 key = node->key;
54 bool leftmost = true;
55
56 while (*new) {
57 parent = *new;
58 if (key < rb_entry(parent, struct test_node, rb)->key)
59 new = &parent->rb_left;
60 else {
61 new = &parent->rb_right;
62 leftmost = false;
63 }
64 }
65
66 rb_link_node(node: &node->rb, parent, rb_link: new);
67 rb_insert_color_cached(node: &node->rb, root, leftmost);
68}
69
70static inline void erase(struct test_node *node, struct rb_root_cached *root)
71{
72 rb_erase(&node->rb, &root->rb_root);
73}
74
75static inline void erase_cached(struct test_node *node, struct rb_root_cached *root)
76{
77 rb_erase_cached(node: &node->rb, root);
78}
79
80
81#define NODE_VAL(node) ((node)->val)
82
83RB_DECLARE_CALLBACKS_MAX(static, augment_callbacks,
84 struct test_node, rb, u32, augmented, NODE_VAL)
85
86static void insert_augmented(struct test_node *node,
87 struct rb_root_cached *root)
88{
89 struct rb_node **new = &root->rb_root.rb_node, *rb_parent = NULL;
90 u32 key = node->key;
91 u32 val = node->val;
92 struct test_node *parent;
93
94 while (*new) {
95 rb_parent = *new;
96 parent = rb_entry(rb_parent, struct test_node, rb);
97 if (parent->augmented < val)
98 parent->augmented = val;
99 if (key < parent->key)
100 new = &parent->rb.rb_left;
101 else
102 new = &parent->rb.rb_right;
103 }
104
105 node->augmented = val;
106 rb_link_node(node: &node->rb, parent: rb_parent, rb_link: new);
107 rb_insert_augmented(node: &node->rb, root: &root->rb_root, augment: &augment_callbacks);
108}
109
110static void insert_augmented_cached(struct test_node *node,
111 struct rb_root_cached *root)
112{
113 struct rb_node **new = &root->rb_root.rb_node, *rb_parent = NULL;
114 u32 key = node->key;
115 u32 val = node->val;
116 struct test_node *parent;
117 bool leftmost = true;
118
119 while (*new) {
120 rb_parent = *new;
121 parent = rb_entry(rb_parent, struct test_node, rb);
122 if (parent->augmented < val)
123 parent->augmented = val;
124 if (key < parent->key)
125 new = &parent->rb.rb_left;
126 else {
127 new = &parent->rb.rb_right;
128 leftmost = false;
129 }
130 }
131
132 node->augmented = val;
133 rb_link_node(node: &node->rb, parent: rb_parent, rb_link: new);
134 rb_insert_augmented_cached(node: &node->rb, root,
135 newleft: leftmost, augment: &augment_callbacks);
136}
137
138
139static void erase_augmented(struct test_node *node, struct rb_root_cached *root)
140{
141 rb_erase_augmented(node: &node->rb, root: &root->rb_root, augment: &augment_callbacks);
142}
143
144static void erase_augmented_cached(struct test_node *node,
145 struct rb_root_cached *root)
146{
147 rb_erase_augmented_cached(node: &node->rb, root, augment: &augment_callbacks);
148}
149
150static void init(void)
151{
152 int i;
153 for (i = 0; i < nnodes; i++) {
154 nodes[i].key = prandom_u32_state(state: &rnd);
155 nodes[i].val = prandom_u32_state(state: &rnd);
156 }
157}
158
159static bool is_red(struct rb_node *rb)
160{
161 return !(rb->__rb_parent_color & 1);
162}
163
164static int black_path_count(struct rb_node *rb)
165{
166 int count;
167 for (count = 0; rb; rb = rb_parent(rb))
168 count += !is_red(rb);
169 return count;
170}
171
172static void check_postorder_foreach(int nr_nodes)
173{
174 struct test_node *cur, *n;
175 int count = 0;
176 rbtree_postorder_for_each_entry_safe(cur, n, &root.rb_root, rb)
177 count++;
178
179 WARN_ON_ONCE(count != nr_nodes);
180}
181
182static void check_postorder(int nr_nodes)
183{
184 struct rb_node *rb;
185 int count = 0;
186 for (rb = rb_first_postorder(&root.rb_root); rb; rb = rb_next_postorder(rb))
187 count++;
188
189 WARN_ON_ONCE(count != nr_nodes);
190}
191
192static void check(int nr_nodes)
193{
194 struct rb_node *rb;
195 int count = 0, blacks = 0;
196 u32 prev_key = 0;
197
198 for (rb = rb_first(root: &root.rb_root); rb; rb = rb_next(rb)) {
199 struct test_node *node = rb_entry(rb, struct test_node, rb);
200 WARN_ON_ONCE(node->key < prev_key);
201 WARN_ON_ONCE(is_red(rb) &&
202 (!rb_parent(rb) || is_red(rb_parent(rb))));
203 if (!count)
204 blacks = black_path_count(rb);
205 else
206 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
207 blacks != black_path_count(rb));
208 prev_key = node->key;
209 count++;
210 }
211
212 WARN_ON_ONCE(count != nr_nodes);
213 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root.rb_root))) - 1);
214
215 check_postorder(nr_nodes);
216 check_postorder_foreach(nr_nodes);
217}
218
219static void check_augmented(int nr_nodes)
220{
221 struct rb_node *rb;
222
223 check(nr_nodes);
224 for (rb = rb_first(root: &root.rb_root); rb; rb = rb_next(rb)) {
225 struct test_node *node = rb_entry(rb, struct test_node, rb);
226 u32 subtree, max = node->val;
227 if (node->rb.rb_left) {
228 subtree = rb_entry(node->rb.rb_left, struct test_node,
229 rb)->augmented;
230 if (max < subtree)
231 max = subtree;
232 }
233 if (node->rb.rb_right) {
234 subtree = rb_entry(node->rb.rb_right, struct test_node,
235 rb)->augmented;
236 if (max < subtree)
237 max = subtree;
238 }
239 WARN_ON_ONCE(node->augmented != max);
240 }
241}
242
243static int basic_check(void)
244{
245 int i, j;
246 cycles_t time1, time2, time;
247 struct rb_node *node;
248
249 printk(KERN_ALERT "rbtree testing");
250
251 init();
252
253 time1 = get_cycles();
254
255 for (i = 0; i < perf_loops; i++) {
256 for (j = 0; j < nnodes; j++)
257 insert(node: nodes + j, root: &root);
258 for (j = 0; j < nnodes; j++)
259 erase(node: nodes + j, root: &root);
260 }
261
262 time2 = get_cycles();
263 time = time2 - time1;
264
265 time = div_u64(dividend: time, divisor: perf_loops);
266 printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n",
267 (unsigned long long)time);
268
269 time1 = get_cycles();
270
271 for (i = 0; i < perf_loops; i++) {
272 for (j = 0; j < nnodes; j++)
273 insert_cached(node: nodes + j, root: &root);
274 for (j = 0; j < nnodes; j++)
275 erase_cached(node: nodes + j, root: &root);
276 }
277
278 time2 = get_cycles();
279 time = time2 - time1;
280
281 time = div_u64(dividend: time, divisor: perf_loops);
282 printk(" -> test 2 (latency of nnodes cached insert+delete): %llu cycles\n",
283 (unsigned long long)time);
284
285 for (i = 0; i < nnodes; i++)
286 insert(node: nodes + i, root: &root);
287
288 time1 = get_cycles();
289
290 for (i = 0; i < perf_loops; i++) {
291 for (node = rb_first(root: &root.rb_root); node; node = rb_next(node))
292 ;
293 }
294
295 time2 = get_cycles();
296 time = time2 - time1;
297
298 time = div_u64(dividend: time, divisor: perf_loops);
299 printk(" -> test 3 (latency of inorder traversal): %llu cycles\n",
300 (unsigned long long)time);
301
302 time1 = get_cycles();
303
304 for (i = 0; i < perf_loops; i++)
305 node = rb_first(root: &root.rb_root);
306
307 time2 = get_cycles();
308 time = time2 - time1;
309
310 time = div_u64(dividend: time, divisor: perf_loops);
311 printk(" -> test 4 (latency to fetch first node)\n");
312 printk(" non-cached: %llu cycles\n", (unsigned long long)time);
313
314 time1 = get_cycles();
315
316 for (i = 0; i < perf_loops; i++)
317 node = rb_first_cached(&root);
318
319 time2 = get_cycles();
320 time = time2 - time1;
321
322 time = div_u64(dividend: time, divisor: perf_loops);
323 printk(" cached: %llu cycles\n", (unsigned long long)time);
324
325 for (i = 0; i < nnodes; i++)
326 erase(node: nodes + i, root: &root);
327
328 /* run checks */
329 for (i = 0; i < check_loops; i++) {
330 init();
331 for (j = 0; j < nnodes; j++) {
332 check(nr_nodes: j);
333 insert(node: nodes + j, root: &root);
334 }
335 for (j = 0; j < nnodes; j++) {
336 check(nr_nodes: nnodes - j);
337 erase(node: nodes + j, root: &root);
338 }
339 check(nr_nodes: 0);
340 }
341
342 return 0;
343}
344
345static int augmented_check(void)
346{
347 int i, j;
348 cycles_t time1, time2, time;
349
350 printk(KERN_ALERT "augmented rbtree testing");
351
352 init();
353
354 time1 = get_cycles();
355
356 for (i = 0; i < perf_loops; i++) {
357 for (j = 0; j < nnodes; j++)
358 insert_augmented(node: nodes + j, root: &root);
359 for (j = 0; j < nnodes; j++)
360 erase_augmented(node: nodes + j, root: &root);
361 }
362
363 time2 = get_cycles();
364 time = time2 - time1;
365
366 time = div_u64(dividend: time, divisor: perf_loops);
367 printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n", (unsigned long long)time);
368
369 time1 = get_cycles();
370
371 for (i = 0; i < perf_loops; i++) {
372 for (j = 0; j < nnodes; j++)
373 insert_augmented_cached(node: nodes + j, root: &root);
374 for (j = 0; j < nnodes; j++)
375 erase_augmented_cached(node: nodes + j, root: &root);
376 }
377
378 time2 = get_cycles();
379 time = time2 - time1;
380
381 time = div_u64(dividend: time, divisor: perf_loops);
382 printk(" -> test 2 (latency of nnodes cached insert+delete): %llu cycles\n", (unsigned long long)time);
383
384 for (i = 0; i < check_loops; i++) {
385 init();
386 for (j = 0; j < nnodes; j++) {
387 check_augmented(nr_nodes: j);
388 insert_augmented(node: nodes + j, root: &root);
389 }
390 for (j = 0; j < nnodes; j++) {
391 check_augmented(nr_nodes: nnodes - j);
392 erase_augmented(node: nodes + j, root: &root);
393 }
394 check_augmented(nr_nodes: 0);
395 }
396
397 return 0;
398}
399
400static int __init rbtree_test_init(void)
401{
402 nodes = kmalloc_array(nnodes, sizeof(*nodes), GFP_KERNEL);
403 if (!nodes)
404 return -ENOMEM;
405
406 prandom_seed_state(state: &rnd, seed);
407
408 basic_check();
409 augmented_check();
410
411 kfree(objp: nodes);
412
413 return -EAGAIN; /* Fail will directly unload the module */
414}
415
416static void __exit rbtree_test_exit(void)
417{
418 printk(KERN_ALERT "test exit\n");
419}
420
421module_init(rbtree_test_init)
422module_exit(rbtree_test_exit)
423
424MODULE_LICENSE("GPL");
425MODULE_AUTHOR("Michel Lespinasse");
426MODULE_DESCRIPTION("Red Black Tree test");
427

source code of linux/lib/rbtree_test.c