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