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 struct rb_node rb; 12 u32 key; 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(int nr_nodes) 118 { 119 struct rb_node *rb; 120 int count = 0; 121 for (rb = rb_first_postorder(&root); rb; rb = rb_next_postorder(rb)) 122 count++; 123 124 WARN_ON_ONCE(count != nr_nodes); 125 } 126 127 static void check(int nr_nodes) 128 { 129 struct rb_node *rb; 130 int count = 0, blacks = 0; 131 u32 prev_key = 0; 132 133 for (rb = rb_first(&root); rb; rb = rb_next(rb)) { 134 struct test_node *node = rb_entry(rb, struct test_node, rb); 135 WARN_ON_ONCE(node->key < prev_key); 136 WARN_ON_ONCE(is_red(rb) && 137 (!rb_parent(rb) || is_red(rb_parent(rb)))); 138 if (!count) 139 blacks = black_path_count(rb); 140 else 141 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) && 142 blacks != black_path_count(rb)); 143 prev_key = node->key; 144 count++; 145 } 146 147 WARN_ON_ONCE(count != nr_nodes); 148 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1); 149 150 check_postorder(nr_nodes); 151 } 152 153 static void check_augmented(int nr_nodes) 154 { 155 struct rb_node *rb; 156 157 check(nr_nodes); 158 for (rb = rb_first(&root); rb; rb = rb_next(rb)) { 159 struct test_node *node = rb_entry(rb, struct test_node, rb); 160 WARN_ON_ONCE(node->augmented != augment_recompute(node)); 161 } 162 } 163 164 static int __init rbtree_test_init(void) 165 { 166 int i, j; 167 cycles_t time1, time2, time; 168 169 printk(KERN_ALERT "rbtree testing"); 170 171 prandom_seed_state(&rnd, 3141592653589793238ULL); 172 init(); 173 174 time1 = get_cycles(); 175 176 for (i = 0; i < PERF_LOOPS; i++) { 177 for (j = 0; j < NODES; j++) 178 insert(nodes + j, &root); 179 for (j = 0; j < NODES; j++) 180 erase(nodes + j, &root); 181 } 182 183 time2 = get_cycles(); 184 time = time2 - time1; 185 186 time = div_u64(time, PERF_LOOPS); 187 printk(" -> %llu cycles\n", (unsigned long long)time); 188 189 for (i = 0; i < CHECK_LOOPS; i++) { 190 init(); 191 for (j = 0; j < NODES; j++) { 192 check(j); 193 insert(nodes + j, &root); 194 } 195 for (j = 0; j < NODES; j++) { 196 check(NODES - j); 197 erase(nodes + j, &root); 198 } 199 check(0); 200 } 201 202 printk(KERN_ALERT "augmented rbtree testing"); 203 204 init(); 205 206 time1 = get_cycles(); 207 208 for (i = 0; i < PERF_LOOPS; i++) { 209 for (j = 0; j < NODES; j++) 210 insert_augmented(nodes + j, &root); 211 for (j = 0; j < NODES; j++) 212 erase_augmented(nodes + j, &root); 213 } 214 215 time2 = get_cycles(); 216 time = time2 - time1; 217 218 time = div_u64(time, PERF_LOOPS); 219 printk(" -> %llu cycles\n", (unsigned long long)time); 220 221 for (i = 0; i < CHECK_LOOPS; i++) { 222 init(); 223 for (j = 0; j < NODES; j++) { 224 check_augmented(j); 225 insert_augmented(nodes + j, &root); 226 } 227 for (j = 0; j < NODES; j++) { 228 check_augmented(NODES - j); 229 erase_augmented(nodes + j, &root); 230 } 231 check_augmented(0); 232 } 233 234 return -EAGAIN; /* Fail will directly unload the module */ 235 } 236 237 static void __exit rbtree_test_exit(void) 238 { 239 printk(KERN_ALERT "test exit\n"); 240 } 241 242 module_init(rbtree_test_init) 243 module_exit(rbtree_test_exit) 244 245 MODULE_LICENSE("GPL"); 246 MODULE_AUTHOR("Michel Lespinasse"); 247 MODULE_DESCRIPTION("Red Black Tree test"); 248