blob: 268b23951fec871d77a790c290637b3d3caf75e8 [file] [log] [blame]
Michel Lespinasse910a7422012-10-08 16:30:39 -07001#include <linux/module.h>
Michel Lespinasse9c079ad2012-10-08 16:31:33 -07002#include <linux/rbtree_augmented.h>
Michel Lespinasse910a7422012-10-08 16:30:39 -07003#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;
Michel Lespinassedadf9352012-10-08 16:31:15 -070013
14 /* following fields used for testing augmented rbtree functionality */
15 u32 val;
16 u32 augmented;
Michel Lespinasse910a7422012-10-08 16:30:39 -070017};
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;
Michel Lespinassedadf9352012-10-08 16:31:15 -070027 u32 key = node->key;
Michel Lespinasse910a7422012-10-08 16:30:39 -070028
29 while (*new) {
30 parent = *new;
Michel Lespinassedadf9352012-10-08 16:31:15 -070031 if (key < rb_entry(parent, struct test_node, rb)->key)
Michel Lespinasse910a7422012-10-08 16:30:39 -070032 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
Michel Lespinassedadf9352012-10-08 16:31:15 -070046static 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
Michel Lespinasse39088362012-10-08 16:31:21 -070064RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
65 u32, augmented, augment_recompute)
Michel Lespinasse14b94af2012-10-08 16:31:17 -070066
Michel Lespinassedadf9352012-10-08 16:31:15 -070067static void insert_augmented(struct test_node *node, struct rb_root *root)
68{
Michel Lespinasse14b94af2012-10-08 16:31:17 -070069 struct rb_node **new = &root->rb_node, *rb_parent = NULL;
Michel Lespinassedadf9352012-10-08 16:31:15 -070070 u32 key = node->key;
Michel Lespinasse14b94af2012-10-08 16:31:17 -070071 u32 val = node->val;
72 struct test_node *parent;
Michel Lespinassedadf9352012-10-08 16:31:15 -070073
74 while (*new) {
Michel Lespinasse14b94af2012-10-08 16:31:17 -070075 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;
Michel Lespinassedadf9352012-10-08 16:31:15 -070081 else
Michel Lespinasse14b94af2012-10-08 16:31:17 -070082 new = &parent->rb.rb_right;
Michel Lespinassedadf9352012-10-08 16:31:15 -070083 }
84
Michel Lespinasse14b94af2012-10-08 16:31:17 -070085 node->augmented = val;
86 rb_link_node(&node->rb, rb_parent, new);
87 rb_insert_augmented(&node->rb, root, &augment_callbacks);
Michel Lespinassedadf9352012-10-08 16:31:15 -070088}
89
90static void erase_augmented(struct test_node *node, struct rb_root *root)
91{
Michel Lespinasse14b94af2012-10-08 16:31:17 -070092 rb_erase_augmented(&node->rb, root, &augment_callbacks);
Michel Lespinassedadf9352012-10-08 16:31:15 -070093}
94
Michel Lespinasse910a7422012-10-08 16:30:39 -070095static void init(void)
96{
97 int i;
Michel Lespinassedadf9352012-10-08 16:31:15 -070098 for (i = 0; i < NODES; i++) {
Michel Lespinasse910a7422012-10-08 16:30:39 -070099 nodes[i].key = prandom32(&rnd);
Michel Lespinassedadf9352012-10-08 16:31:15 -0700100 nodes[i].val = prandom32(&rnd);
101 }
Michel Lespinasse910a7422012-10-08 16:30:39 -0700102}
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;
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
Michel Lespinassedadf9352012-10-08 16:31:15 -0700140static 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
Michel Lespinasse910a7422012-10-08 16:30:39 -0700151static int rbtree_test_init(void)
152{
153 int i, j;
154 cycles_t time1, time2, time;
155
156 printk(KERN_ALERT "rbtree testing");
157
Michel Lespinasse28d75302012-10-08 16:31:04 -0700158 prandom32_seed(&rnd, 3141592653589793238ULL);
Michel Lespinasse910a7422012-10-08 16:30:39 -0700159 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
Michel Lespinassedadf9352012-10-08 16:31:15 -0700189 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
Michel Lespinasse910a7422012-10-08 16:30:39 -0700221 return -EAGAIN; /* Fail will directly unload the module */
222}
223
224static void rbtree_test_exit(void)
225{
226 printk(KERN_ALERT "test exit\n");
227}
228
229module_init(rbtree_test_init)
230module_exit(rbtree_test_exit)
231
232MODULE_LICENSE("GPL");
233MODULE_AUTHOR("Michel Lespinasse");
234MODULE_DESCRIPTION("Red Black Tree test");