1 #include <linux/module.h> 2 #include <linux/rbtree_augmented.h> 3 #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 10 struct test_node { 11 u32 key; 12 struct rb_node rb; 13 14 /* following fields used for testing augmented rbtree functionality */ 15 u32 val; 16 u32 augmented; 17 }; 18 19 static struct rb_root root = RB_ROOT; 20 static struct test_node nodes[NODES]; 21 22 static struct rnd_state rnd; 23 24 static void insert(struct test_node *node, struct rb_root *root) 25 { 26 struct rb_node **new = &root->rb_node, *parent = NULL; 27 u32 key = node->key; 28 29 while (*new) { 30 parent = *new; 31 if (key < rb_entry(parent, struct test_node, rb)->key) 32 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 41 static inline void erase(struct test_node *node, struct rb_root *root) 42 { 43 rb_erase(&node->rb, root); 44 } 45 46 static 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 64 RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb, 65 u32, augmented, augment_recompute) 66 67 static void insert_augmented(struct test_node *node, struct rb_root *root) 68 { 69 struct rb_node **new = &root->rb_node, *rb_parent = NULL; 70 u32 key = node->key; 71 u32 val = node->val; 72 struct test_node *parent; 73 74 while (*new) { 75 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; 81 else 82 new = &parent->rb.rb_right; 83 } 84 85 node->augmented = val; 86 rb_link_node(&node->rb, rb_parent, new); 87 rb_insert_augmented(&node->rb, root, &augment_callbacks); 88 } 89 90 static void erase_augmented(struct test_node *node, struct rb_root *root) 91 { 92 rb_erase_augmented(&node->rb, root, &augment_callbacks); 93 } 94 95 static void init(void) 96 { 97 int i; 98 for (i = 0; i < NODES; i++) { 99 nodes[i].key = prandom_u32_state(&rnd); 100 nodes[i].val = prandom_u32_state(&rnd); 101 } 102 } 103 104 static bool is_red(struct rb_node *rb) 105 { 106 return !(rb->__rb_parent_color & 1); 107 } 108 109 static 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 117 static void check_postorder_foreach(int nr_nodes) 118 { 119 struct test_node *cur, *n; 120 int count = 0; 121 rbtree_postorder_for_each_entry_safe(cur, n, &root, rb) 122 count++; 123 124 WARN_ON_ONCE(count != nr_nodes); 125 } 126 127 static void check_postorder(int nr_nodes) 128 { 129 struct rb_node *rb; 130 int count = 0; 131 for (rb = rb_first_postorder(&root); rb; rb = rb_next_postorder(rb)) 132 count++; 133 134 WARN_ON_ONCE(count != nr_nodes); 135 } 136 137 static void check(int nr_nodes) 138 { 139 struct rb_node *rb; 140 int count = 0, blacks = 0; 141 u32 prev_key = 0; 142 143 for (rb = rb_first(&root); rb; rb = rb_next(rb)) { 144 struct test_node *node = rb_entry(rb, struct test_node, rb); 145 WARN_ON_ONCE(node->key < prev_key); 146 WARN_ON_ONCE(is_red(rb) && 147 (!rb_parent(rb) || is_red(rb_parent(rb)))); 148 if (!count) 149 blacks = black_path_count(rb); 150 else 151 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) && 152 blacks != black_path_count(rb)); 153 prev_key = node->key; 154 count++; 155 } 156 157 WARN_ON_ONCE(count != nr_nodes); 158 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1); 159 160 check_postorder(nr_nodes); 161 check_postorder_foreach(nr_nodes); 162 } 163 164 static void check_augmented(int nr_nodes) 165 { 166 struct rb_node *rb; 167 168 check(nr_nodes); 169 for (rb = rb_first(&root); rb; rb = rb_next(rb)) { 170 struct test_node *node = rb_entry(rb, struct test_node, rb); 171 WARN_ON_ONCE(node->augmented != augment_recompute(node)); 172 } 173 } 174 175 static int __init rbtree_test_init(void) 176 { 177 int i, j; 178 cycles_t time1, time2, time; 179 180 printk(KERN_ALERT "rbtree testing"); 181 182 prandom_seed_state(&rnd, 3141592653589793238ULL); 183 init(); 184 185 time1 = get_cycles(); 186 187 for (i = 0; i < PERF_LOOPS; i++) { 188 for (j = 0; j < NODES; j++) 189 insert(nodes + j, &root); 190 for (j = 0; j < NODES; j++) 191 erase(nodes + j, &root); 192 } 193 194 time2 = get_cycles(); 195 time = time2 - time1; 196 197 time = div_u64(time, PERF_LOOPS); 198 printk(" -> %llu cycles\n", (unsigned long long)time); 199 200 for (i = 0; i < CHECK_LOOPS; i++) { 201 init(); 202 for (j = 0; j < NODES; j++) { 203 check(j); 204 insert(nodes + j, &root); 205 } 206 for (j = 0; j < NODES; j++) { 207 check(NODES - j); 208 erase(nodes + j, &root); 209 } 210 check(0); 211 } 212 213 printk(KERN_ALERT "augmented rbtree testing"); 214 215 init(); 216 217 time1 = get_cycles(); 218 219 for (i = 0; i < PERF_LOOPS; i++) { 220 for (j = 0; j < NODES; j++) 221 insert_augmented(nodes + j, &root); 222 for (j = 0; j < NODES; j++) 223 erase_augmented(nodes + j, &root); 224 } 225 226 time2 = get_cycles(); 227 time = time2 - time1; 228 229 time = div_u64(time, PERF_LOOPS); 230 printk(" -> %llu cycles\n", (unsigned long long)time); 231 232 for (i = 0; i < CHECK_LOOPS; i++) { 233 init(); 234 for (j = 0; j < NODES; j++) { 235 check_augmented(j); 236 insert_augmented(nodes + j, &root); 237 } 238 for (j = 0; j < NODES; j++) { 239 check_augmented(NODES - j); 240 erase_augmented(nodes + j, &root); 241 } 242 check_augmented(0); 243 } 244 245 return -EAGAIN; /* Fail will directly unload the module */ 246 } 247 248 static void __exit rbtree_test_exit(void) 249 { 250 printk(KERN_ALERT "test exit\n"); 251 } 252 253 module_init(rbtree_test_init) 254 module_exit(rbtree_test_exit) 255 256 MODULE_LICENSE("GPL"); 257 MODULE_AUTHOR("Michel Lespinasse"); 258 MODULE_DESCRIPTION("Red Black Tree test"); 259