1 #include <linux/module.h>
2 #include <linux/rbtree.h>
3 #include <linux/random.h>
7 #define PERF_LOOPS 100000
8 #define CHECK_LOOPS 100
14 /* following fields used for testing augmented rbtree functionality */
19 static struct rb_root root = RB_ROOT;
20 static struct test_node nodes[NODES];
22 static struct rnd_state rnd;
24 static void insert(struct test_node *node, struct rb_root *root)
26 struct rb_node **new = &root->rb_node, *parent = NULL;
31 if (key < rb_entry(parent, struct test_node, rb)->key)
32 new = &parent->rb_left;
34 new = &parent->rb_right;
37 rb_link_node(&node->rb, parent, new);
38 rb_insert_color(&node->rb, root);
41 static inline void erase(struct test_node *node, struct rb_root *root)
43 rb_erase(&node->rb, root);
46 static inline u32 augment_recompute(struct test_node *node)
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,
52 if (max < child_augmented)
53 max = child_augmented;
55 if (node->rb.rb_right) {
56 child_augmented = rb_entry(node->rb.rb_right, struct test_node,
58 if (max < child_augmented)
59 max = child_augmented;
64 static void augment_propagate(struct rb_node *rb, struct rb_node *stop)
67 struct test_node *node = rb_entry(rb, struct test_node, rb);
68 u32 augmented = augment_recompute(node);
69 if (node->augmented == augmented)
71 node->augmented = augmented;
72 rb = rb_parent(&node->rb);
76 static void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new)
78 struct test_node *old = rb_entry(rb_old, struct test_node, rb);
79 struct test_node *new = rb_entry(rb_new, struct test_node, rb);
80 new->augmented = old->augmented;
83 static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
85 struct test_node *old = rb_entry(rb_old, struct test_node, rb);
86 struct test_node *new = rb_entry(rb_new, struct test_node, rb);
88 /* Rotation doesn't change subtree's augmented value */
89 new->augmented = old->augmented;
90 old->augmented = augment_recompute(old);
93 static const struct rb_augment_callbacks augment_callbacks = {
94 augment_propagate, augment_copy, augment_rotate
97 static void insert_augmented(struct test_node *node, struct rb_root *root)
99 struct rb_node **new = &root->rb_node, *rb_parent = NULL;
102 struct test_node *parent;
106 parent = rb_entry(rb_parent, struct test_node, rb);
107 if (parent->augmented < val)
108 parent->augmented = val;
109 if (key < parent->key)
110 new = &parent->rb.rb_left;
112 new = &parent->rb.rb_right;
115 node->augmented = val;
116 rb_link_node(&node->rb, rb_parent, new);
117 rb_insert_augmented(&node->rb, root, &augment_callbacks);
120 static void erase_augmented(struct test_node *node, struct rb_root *root)
122 rb_erase_augmented(&node->rb, root, &augment_callbacks);
125 static void init(void)
128 for (i = 0; i < NODES; i++) {
129 nodes[i].key = prandom32(&rnd);
130 nodes[i].val = prandom32(&rnd);
134 static bool is_red(struct rb_node *rb)
136 return !(rb->__rb_parent_color & 1);
139 static int black_path_count(struct rb_node *rb)
142 for (count = 0; rb; rb = rb_parent(rb))
143 count += !is_red(rb);
147 static void check(int nr_nodes)
154 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
155 struct test_node *node = rb_entry(rb, struct test_node, rb);
156 WARN_ON_ONCE(node->key < prev_key);
157 WARN_ON_ONCE(is_red(rb) &&
158 (!rb_parent(rb) || is_red(rb_parent(rb))));
160 blacks = black_path_count(rb);
162 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
163 blacks != black_path_count(rb));
164 prev_key = node->key;
167 WARN_ON_ONCE(count != nr_nodes);
170 static void check_augmented(int nr_nodes)
175 for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
176 struct test_node *node = rb_entry(rb, struct test_node, rb);
177 WARN_ON_ONCE(node->augmented != augment_recompute(node));
181 static int rbtree_test_init(void)
184 cycles_t time1, time2, time;
186 printk(KERN_ALERT "rbtree testing");
188 prandom32_seed(&rnd, 3141592653589793238ULL);
191 time1 = get_cycles();
193 for (i = 0; i < PERF_LOOPS; i++) {
194 for (j = 0; j < NODES; j++)
195 insert(nodes + j, &root);
196 for (j = 0; j < NODES; j++)
197 erase(nodes + j, &root);
200 time2 = get_cycles();
201 time = time2 - time1;
203 time = div_u64(time, PERF_LOOPS);
204 printk(" -> %llu cycles\n", (unsigned long long)time);
206 for (i = 0; i < CHECK_LOOPS; i++) {
208 for (j = 0; j < NODES; j++) {
210 insert(nodes + j, &root);
212 for (j = 0; j < NODES; j++) {
214 erase(nodes + j, &root);
219 printk(KERN_ALERT "augmented rbtree testing");
223 time1 = get_cycles();
225 for (i = 0; i < PERF_LOOPS; i++) {
226 for (j = 0; j < NODES; j++)
227 insert_augmented(nodes + j, &root);
228 for (j = 0; j < NODES; j++)
229 erase_augmented(nodes + j, &root);
232 time2 = get_cycles();
233 time = time2 - time1;
235 time = div_u64(time, PERF_LOOPS);
236 printk(" -> %llu cycles\n", (unsigned long long)time);
238 for (i = 0; i < CHECK_LOOPS; i++) {
240 for (j = 0; j < NODES; j++) {
242 insert_augmented(nodes + j, &root);
244 for (j = 0; j < NODES; j++) {
245 check_augmented(NODES - j);
246 erase_augmented(nodes + j, &root);
251 return -EAGAIN; /* Fail will directly unload the module */
254 static void rbtree_test_exit(void)
256 printk(KERN_ALERT "test exit\n");
259 module_init(rbtree_test_init)
260 module_exit(rbtree_test_exit)
262 MODULE_LICENSE("GPL");
263 MODULE_AUTHOR("Michel Lespinasse");
264 MODULE_DESCRIPTION("Red Black Tree test");