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
10struct test_node {
11 struct rb_node rb;
12 u32 key;
13
14
15 u32 val;
16 u32 augmented;
17};
18
19static struct rb_root root = RB_ROOT;
20static struct test_node nodes[NODES];
21
22static struct rnd_state rnd;
23
24static 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
41static inline void erase(struct test_node *node, struct rb_root *root)
42{
43 rb_erase(&node->rb, root);
44}
45
46static 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
64RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
65 u32, augmented, augment_recompute)
66
67static 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
90static void erase_augmented(struct test_node *node, struct rb_root *root)
91{
92 rb_erase_augmented(&node->rb, root, &augment_callbacks);
93}
94
95static 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
104static bool is_red(struct rb_node *rb)
105{
106 return !(rb->__rb_parent_color & 1);
107}
108
109static 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
117static 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
141static 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
152static 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;
223}
224
225static void __exit rbtree_test_exit(void)
226{
227 printk(KERN_ALERT "test exit\n");
228}
229
230module_init(rbtree_test_init)
231module_exit(rbtree_test_exit)
232
233MODULE_LICENSE("GPL");
234MODULE_AUTHOR("Michel Lespinasse");
235MODULE_DESCRIPTION("Red Black Tree test");
236
