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