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