Merge branch 'parisc-3.11-rc4' of git://git.kernel.org/pub/scm/linux/kernel/git/delle...
[linux-2.6.git] / lib / rbtree_test.c
blob122f02f9941b5680b6fd9f14793e8806476bb634
1 #include <linux/module.h>
2 #include <linux/rbtree_augmented.h>
3 #include <linux/random.h>
4 #include <asm/timex.h>
6 #define NODES 100
7 #define PERF_LOOPS 100000
8 #define CHECK_LOOPS 100
10 struct test_node {
11 struct rb_node rb;
12 u32 key;
14 /* following fields used for testing augmented rbtree functionality */
15 u32 val;
16 u32 augmented;
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;
27 u32 key = node->key;
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;
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,
51 rb)->augmented;
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,
57 rb)->augmented;
58 if (max < child_augmented)
59 max = child_augmented;
61 return max;
64 RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
65 u32, augmented, augment_recompute)
67 static void insert_augmented(struct test_node *node, struct rb_root *root)
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;
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;
85 node->augmented = val;
86 rb_link_node(&node->rb, rb_parent, new);
87 rb_insert_augmented(&node->rb, root, &augment_callbacks);
90 static void erase_augmented(struct test_node *node, struct rb_root *root)
92 rb_erase_augmented(&node->rb, root, &augment_callbacks);
95 static void init(void)
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);
104 static bool is_red(struct rb_node *rb)
106 return !(rb->__rb_parent_color & 1);
109 static int black_path_count(struct rb_node *rb)
111 int count;
112 for (count = 0; rb; rb = rb_parent(rb))
113 count += !is_red(rb);
114 return count;
117 static void check(int nr_nodes)
119 struct rb_node *rb;
120 int count = 0, blacks = 0;
121 u32 prev_key = 0;
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++;
137 WARN_ON_ONCE(count != nr_nodes);
138 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);
141 static void check_augmented(int nr_nodes)
143 struct rb_node *rb;
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));
152 static int __init rbtree_test_init(void)
154 int i, j;
155 cycles_t time1, time2, time;
157 printk(KERN_ALERT "rbtree testing");
159 prandom_seed_state(&rnd, 3141592653589793238ULL);
160 init();
162 time1 = get_cycles();
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);
171 time2 = get_cycles();
172 time = time2 - time1;
174 time = div_u64(time, PERF_LOOPS);
175 printk(" -> %llu cycles\n", (unsigned long long)time);
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);
183 for (j = 0; j < NODES; j++) {
184 check(NODES - j);
185 erase(nodes + j, &root);
187 check(0);
190 printk(KERN_ALERT "augmented rbtree testing");
192 init();
194 time1 = get_cycles();
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);
203 time2 = get_cycles();
204 time = time2 - time1;
206 time = div_u64(time, PERF_LOOPS);
207 printk(" -> %llu cycles\n", (unsigned long long)time);
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);
215 for (j = 0; j < NODES; j++) {
216 check_augmented(NODES - j);
217 erase_augmented(nodes + j, &root);
219 check_augmented(0);
222 return -EAGAIN; /* Fail will directly unload the module */
225 static void __exit rbtree_test_exit(void)
227 printk(KERN_ALERT "test exit\n");
230 module_init(rbtree_test_init)
231 module_exit(rbtree_test_exit)
233 MODULE_LICENSE("GPL");
234 MODULE_AUTHOR("Michel Lespinasse");
235 MODULE_DESCRIPTION("Red Black Tree test");