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