]>
git.proxmox.com Git - mirror_ubuntu-jammy-kernel.git/blob - lib/rbtree_test.c
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/random.h>
6 #include <linux/slab.h>
9 #define __param(type, name, init, msg) \
10 static type name = init; \
11 module_param(name, type, 0444); \
12 MODULE_PARM_DESC(name, msg);
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");
22 /* following fields used for testing augmented rbtree functionality */
27 static struct rb_root_cached root
= RB_ROOT_CACHED
;
28 static struct test_node
*nodes
= NULL
;
30 static struct rnd_state rnd
;
32 static void insert(struct test_node
*node
, struct rb_root_cached
*root
)
34 struct rb_node
**new = &root
->rb_root
.rb_node
, *parent
= NULL
;
39 if (key
< rb_entry(parent
, struct test_node
, rb
)->key
)
40 new = &parent
->rb_left
;
42 new = &parent
->rb_right
;
45 rb_link_node(&node
->rb
, parent
, new);
46 rb_insert_color(&node
->rb
, &root
->rb_root
);
49 static void insert_cached(struct test_node
*node
, struct rb_root_cached
*root
)
51 struct rb_node
**new = &root
->rb_root
.rb_node
, *parent
= NULL
;
57 if (key
< rb_entry(parent
, struct test_node
, rb
)->key
)
58 new = &parent
->rb_left
;
60 new = &parent
->rb_right
;
65 rb_link_node(&node
->rb
, parent
, new);
66 rb_insert_color_cached(&node
->rb
, root
, leftmost
);
69 static inline void erase(struct test_node
*node
, struct rb_root_cached
*root
)
71 rb_erase(&node
->rb
, &root
->rb_root
);
74 static inline void erase_cached(struct test_node
*node
, struct rb_root_cached
*root
)
76 rb_erase_cached(&node
->rb
, root
);
80 static inline u32
augment_recompute(struct test_node
*node
)
82 u32 max
= node
->val
, child_augmented
;
83 if (node
->rb
.rb_left
) {
84 child_augmented
= rb_entry(node
->rb
.rb_left
, struct test_node
,
86 if (max
< child_augmented
)
87 max
= child_augmented
;
89 if (node
->rb
.rb_right
) {
90 child_augmented
= rb_entry(node
->rb
.rb_right
, struct test_node
,
92 if (max
< child_augmented
)
93 max
= child_augmented
;
98 RB_DECLARE_CALLBACKS(static, augment_callbacks
, struct test_node
, rb
,
99 u32
, augmented
, augment_recompute
)
101 static void insert_augmented(struct test_node
*node
,
102 struct rb_root_cached
*root
)
104 struct rb_node
**new = &root
->rb_root
.rb_node
, *rb_parent
= NULL
;
107 struct test_node
*parent
;
111 parent
= rb_entry(rb_parent
, struct test_node
, rb
);
112 if (parent
->augmented
< val
)
113 parent
->augmented
= val
;
114 if (key
< parent
->key
)
115 new = &parent
->rb
.rb_left
;
117 new = &parent
->rb
.rb_right
;
120 node
->augmented
= val
;
121 rb_link_node(&node
->rb
, rb_parent
, new);
122 rb_insert_augmented(&node
->rb
, &root
->rb_root
, &augment_callbacks
);
125 static void insert_augmented_cached(struct test_node
*node
,
126 struct rb_root_cached
*root
)
128 struct rb_node
**new = &root
->rb_root
.rb_node
, *rb_parent
= NULL
;
131 struct test_node
*parent
;
132 bool leftmost
= true;
136 parent
= rb_entry(rb_parent
, struct test_node
, rb
);
137 if (parent
->augmented
< val
)
138 parent
->augmented
= val
;
139 if (key
< parent
->key
)
140 new = &parent
->rb
.rb_left
;
142 new = &parent
->rb
.rb_right
;
147 node
->augmented
= val
;
148 rb_link_node(&node
->rb
, rb_parent
, new);
149 rb_insert_augmented_cached(&node
->rb
, root
,
150 leftmost
, &augment_callbacks
);
154 static void erase_augmented(struct test_node
*node
, struct rb_root_cached
*root
)
156 rb_erase_augmented(&node
->rb
, &root
->rb_root
, &augment_callbacks
);
159 static void erase_augmented_cached(struct test_node
*node
,
160 struct rb_root_cached
*root
)
162 rb_erase_augmented_cached(&node
->rb
, root
, &augment_callbacks
);
165 static void init(void)
168 for (i
= 0; i
< nnodes
; i
++) {
169 nodes
[i
].key
= prandom_u32_state(&rnd
);
170 nodes
[i
].val
= prandom_u32_state(&rnd
);
174 static bool is_red(struct rb_node
*rb
)
176 return !(rb
->__rb_parent_color
& 1);
179 static int black_path_count(struct rb_node
*rb
)
182 for (count
= 0; rb
; rb
= rb_parent(rb
))
183 count
+= !is_red(rb
);
187 static void check_postorder_foreach(int nr_nodes
)
189 struct test_node
*cur
, *n
;
191 rbtree_postorder_for_each_entry_safe(cur
, n
, &root
.rb_root
, rb
)
194 WARN_ON_ONCE(count
!= nr_nodes
);
197 static void check_postorder(int nr_nodes
)
201 for (rb
= rb_first_postorder(&root
.rb_root
); rb
; rb
= rb_next_postorder(rb
))
204 WARN_ON_ONCE(count
!= nr_nodes
);
207 static void check(int nr_nodes
)
210 int count
= 0, blacks
= 0;
213 for (rb
= rb_first(&root
.rb_root
); rb
; rb
= rb_next(rb
)) {
214 struct test_node
*node
= rb_entry(rb
, struct test_node
, rb
);
215 WARN_ON_ONCE(node
->key
< prev_key
);
216 WARN_ON_ONCE(is_red(rb
) &&
217 (!rb_parent(rb
) || is_red(rb_parent(rb
))));
219 blacks
= black_path_count(rb
);
221 WARN_ON_ONCE((!rb
->rb_left
|| !rb
->rb_right
) &&
222 blacks
!= black_path_count(rb
));
223 prev_key
= node
->key
;
227 WARN_ON_ONCE(count
!= nr_nodes
);
228 WARN_ON_ONCE(count
< (1 << black_path_count(rb_last(&root
.rb_root
))) - 1);
230 check_postorder(nr_nodes
);
231 check_postorder_foreach(nr_nodes
);
234 static void check_augmented(int nr_nodes
)
239 for (rb
= rb_first(&root
.rb_root
); rb
; rb
= rb_next(rb
)) {
240 struct test_node
*node
= rb_entry(rb
, struct test_node
, rb
);
241 WARN_ON_ONCE(node
->augmented
!= augment_recompute(node
));
245 static int __init
rbtree_test_init(void)
248 cycles_t time1
, time2
, time
;
249 struct rb_node
*node
;
251 nodes
= kmalloc_array(nnodes
, sizeof(*nodes
), GFP_KERNEL
);
255 printk(KERN_ALERT
"rbtree testing");
257 prandom_seed_state(&rnd
, 3141592653589793238ULL);
260 time1
= get_cycles();
262 for (i
= 0; i
< perf_loops
; i
++) {
263 for (j
= 0; j
< nnodes
; j
++)
264 insert(nodes
+ j
, &root
);
265 for (j
= 0; j
< nnodes
; j
++)
266 erase(nodes
+ j
, &root
);
269 time2
= get_cycles();
270 time
= time2
- time1
;
272 time
= div_u64(time
, perf_loops
);
273 printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n",
274 (unsigned long long)time
);
276 time1
= get_cycles();
278 for (i
= 0; i
< perf_loops
; i
++) {
279 for (j
= 0; j
< nnodes
; j
++)
280 insert_cached(nodes
+ j
, &root
);
281 for (j
= 0; j
< nnodes
; j
++)
282 erase_cached(nodes
+ j
, &root
);
285 time2
= get_cycles();
286 time
= time2
- time1
;
288 time
= div_u64(time
, perf_loops
);
289 printk(" -> test 2 (latency of nnodes cached insert+delete): %llu cycles\n",
290 (unsigned long long)time
);
292 for (i
= 0; i
< nnodes
; i
++)
293 insert(nodes
+ i
, &root
);
295 time1
= get_cycles();
297 for (i
= 0; i
< perf_loops
; i
++) {
298 for (node
= rb_first(&root
.rb_root
); node
; node
= rb_next(node
))
302 time2
= get_cycles();
303 time
= time2
- time1
;
305 time
= div_u64(time
, perf_loops
);
306 printk(" -> test 3 (latency of inorder traversal): %llu cycles\n",
307 (unsigned long long)time
);
309 time1
= get_cycles();
311 for (i
= 0; i
< perf_loops
; i
++)
312 node
= rb_first(&root
.rb_root
);
314 time2
= get_cycles();
315 time
= time2
- time1
;
317 time
= div_u64(time
, perf_loops
);
318 printk(" -> test 4 (latency to fetch first node)\n");
319 printk(" non-cached: %llu cycles\n", (unsigned long long)time
);
321 time1
= get_cycles();
323 for (i
= 0; i
< perf_loops
; i
++)
324 node
= rb_first_cached(&root
);
326 time2
= get_cycles();
327 time
= time2
- time1
;
329 time
= div_u64(time
, perf_loops
);
330 printk(" cached: %llu cycles\n", (unsigned long long)time
);
332 for (i
= 0; i
< nnodes
; i
++)
333 erase(nodes
+ i
, &root
);
336 for (i
= 0; i
< check_loops
; i
++) {
338 for (j
= 0; j
< nnodes
; j
++) {
340 insert(nodes
+ j
, &root
);
342 for (j
= 0; j
< nnodes
; j
++) {
344 erase(nodes
+ j
, &root
);
349 printk(KERN_ALERT
"augmented rbtree testing");
353 time1
= get_cycles();
355 for (i
= 0; i
< perf_loops
; i
++) {
356 for (j
= 0; j
< nnodes
; j
++)
357 insert_augmented(nodes
+ j
, &root
);
358 for (j
= 0; j
< nnodes
; j
++)
359 erase_augmented(nodes
+ j
, &root
);
362 time2
= get_cycles();
363 time
= time2
- time1
;
365 time
= div_u64(time
, perf_loops
);
366 printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n", (unsigned long long)time
);
368 time1
= get_cycles();
370 for (i
= 0; i
< perf_loops
; i
++) {
371 for (j
= 0; j
< nnodes
; j
++)
372 insert_augmented_cached(nodes
+ j
, &root
);
373 for (j
= 0; j
< nnodes
; j
++)
374 erase_augmented_cached(nodes
+ j
, &root
);
377 time2
= get_cycles();
378 time
= time2
- time1
;
380 time
= div_u64(time
, perf_loops
);
381 printk(" -> test 2 (latency of nnodes cached insert+delete): %llu cycles\n", (unsigned long long)time
);
383 for (i
= 0; i
< check_loops
; i
++) {
385 for (j
= 0; j
< nnodes
; j
++) {
387 insert_augmented(nodes
+ j
, &root
);
389 for (j
= 0; j
< nnodes
; j
++) {
390 check_augmented(nnodes
- j
);
391 erase_augmented(nodes
+ j
, &root
);
398 return -EAGAIN
; /* Fail will directly unload the module */
401 static void __exit
rbtree_test_exit(void)
403 printk(KERN_ALERT
"test exit\n");
406 module_init(rbtree_test_init
)
407 module_exit(rbtree_test_exit
)
409 MODULE_LICENSE("GPL");
410 MODULE_AUTHOR("Michel Lespinasse");
411 MODULE_DESCRIPTION("Red Black Tree test");