projects/linux·blob master · refs

384 lines · 8.6 KiB

1// SPDX-License-Identifier: GPL-2.0-only
2#include <linux/module.h>
3#include <linux/moduleparam.h>
4#include <linux/rbtree_augmented.h>
5#include <linux/prandom.h>
6#include <linux/slab.h>
7#include <asm/timex.h>
8
9#define __param(type, name, init, msg) \
10 static type name = init; \
11 module_param(name, type, 0444); \
12 MODULE_PARM_DESC(name, msg);
13
14__param(int, nnodes, 100, "Number of nodes in the rb-tree");
15__param(int, perf_loops, 1000, "Number of iterations modifying the rb-tree");
16__param(int, check_loops, 100, "Number of iterations modifying and verifying the rb-tree");
17__param(ullong, seed, 3141592653589793238ULL, "Random seed");
18
19struct test_node {
20 u32 key;
21 struct rb_node rb;
22
23 /* following fields used for testing augmented rbtree functionality */
24 u32 val;
25 u32 augmented;
26};
27
28static struct rb_root_cached root = RB_ROOT_CACHED;
29static struct test_node *nodes = NULL;
30
31static struct rnd_state rnd;
32
33static inline bool less(struct rb_node *a, const struct rb_node *b)
34{
35 return rb_entry(a, struct test_node, rb)->key <
36 rb_entry(b, struct test_node, rb)->key;
37}
38
39static void insert(struct test_node *node, struct rb_root_cached *root)
40{
41 rb_add(&node->rb, &root->rb_root, less);
42}
43
44static void insert_cached(struct test_node *node, struct rb_root_cached *root)
45{
46 rb_add_cached(&node->rb, root, less);
47}
48
49static inline void erase(struct test_node *node, struct rb_root_cached *root)
50{
51 rb_erase(&node->rb, &root->rb_root);
52}
53
54static inline void erase_cached(struct test_node *node, struct rb_root_cached *root)
55{
56 rb_erase_cached(&node->rb, root);
57}
58
59
60#define NODE_VAL(node) ((node)->val)
61
62RB_DECLARE_CALLBACKS_MAX(static, augment_callbacks,
63 struct test_node, rb, u32, augmented, NODE_VAL)
64
65static void insert_augmented(struct test_node *node,
66 struct rb_root_cached *root)
67{
68 struct rb_node **new = &root->rb_root.rb_node, *rb_parent = NULL;
69 u32 key = node->key;
70 u32 val = node->val;
71 struct test_node *parent;
72
73 while (*new) {
74 rb_parent = *new;
75 parent = rb_entry(rb_parent, struct test_node, rb);
76 if (parent->augmented < val)
77 parent->augmented = val;
78 if (key < parent->key)
79 new = &parent->rb.rb_left;
80 else
81 new = &parent->rb.rb_right;
82 }
83
84 node->augmented = val;
85 rb_link_node(&node->rb, rb_parent, new);
86 rb_insert_augmented(&node->rb, &root->rb_root, &augment_callbacks);
87}
88
89static void insert_augmented_cached(struct test_node *node,
90 struct rb_root_cached *root)
91{
92 node->augmented = node->val;
93 rb_add_augmented_cached(&node->rb, root, less, &augment_callbacks);
94}
95
96
97static void erase_augmented(struct test_node *node, struct rb_root_cached *root)
98{
99 rb_erase_augmented(&node->rb, &root->rb_root, &augment_callbacks);
100}
101
102static void erase_augmented_cached(struct test_node *node,
103 struct rb_root_cached *root)
104{
105 rb_erase_augmented_cached(&node->rb, root, &augment_callbacks);
106}
107
108static void init(void)
109{
110 int i;
111 for (i = 0; i < nnodes; i++) {
112 nodes[i].key = prandom_u32_state(&rnd);
113 nodes[i].val = prandom_u32_state(&rnd);
114 }
115}
116
117static bool is_red(struct rb_node *rb)
118{
119 return !(rb->__rb_parent_color & 1);
120}
121
122static int black_path_count(struct rb_node *rb)
123{
124 int count;
125 for (count = 0; rb; rb = rb_parent(rb))
126 count += !is_red(rb);
127 return count;
128}
129
130static void check_postorder_foreach(int nr_nodes)
131{
132 struct test_node *cur, *n;
133 int count = 0;
134 rbtree_postorder_for_each_entry_safe(cur, n, &root.rb_root, rb)
135 count++;
136
137 WARN_ON_ONCE(count != nr_nodes);
138}
139
140static void check_postorder(int nr_nodes)
141{
142 struct rb_node *rb;
143 int count = 0;
144 for (rb = rb_first_postorder(&root.rb_root); rb; rb = rb_next_postorder(rb))
145 count++;
146
147 WARN_ON_ONCE(count != nr_nodes);
148}
149
150static void check(int nr_nodes)
151{
152 struct rb_node *rb;
153 int count = 0, blacks = 0;
154 u32 prev_key = 0;
155
156 for (rb = rb_first(&root.rb_root); rb; rb = rb_next(rb)) {
157 struct test_node *node = rb_entry(rb, struct test_node, rb);
158 WARN_ON_ONCE(node->key < prev_key);
159 WARN_ON_ONCE(is_red(rb) &&
160 (!rb_parent(rb) || is_red(rb_parent(rb))));
161 if (!count)
162 blacks = black_path_count(rb);
163 else
164 WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
165 blacks != black_path_count(rb));
166 prev_key = node->key;
167 count++;
168 }
169
170 WARN_ON_ONCE(count != nr_nodes);
171 WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root.rb_root))) - 1);
172
173 check_postorder(nr_nodes);
174 check_postorder_foreach(nr_nodes);
175}
176
177static void check_augmented(int nr_nodes)
178{
179 struct rb_node *rb;
180
181 check(nr_nodes);
182 for (rb = rb_first(&root.rb_root); rb; rb = rb_next(rb)) {
183 struct test_node *node = rb_entry(rb, struct test_node, rb);
184 u32 subtree, max = node->val;
185 if (node->rb.rb_left) {
186 subtree = rb_entry(node->rb.rb_left, struct test_node,
187 rb)->augmented;
188 if (max < subtree)
189 max = subtree;
190 }
191 if (node->rb.rb_right) {
192 subtree = rb_entry(node->rb.rb_right, struct test_node,
193 rb)->augmented;
194 if (max < subtree)
195 max = subtree;
196 }
197 WARN_ON_ONCE(node->augmented != max);
198 }
199}
200
201static int basic_check(void)
202{
203 int i, j;
204 cycles_t time1, time2, time;
205 struct rb_node *node;
206
207 printk(KERN_ALERT "rbtree testing");
208
209 init();
210
211 time1 = get_cycles();
212
213 for (i = 0; i < perf_loops; i++) {
214 for (j = 0; j < nnodes; j++)
215 insert(nodes + j, &root);
216 for (j = 0; j < nnodes; j++)
217 erase(nodes + j, &root);
218 }
219
220 time2 = get_cycles();
221 time = time2 - time1;
222
223 time = div_u64(time, perf_loops);
224 printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n",
225 (unsigned long long)time);
226
227 time1 = get_cycles();
228
229 for (i = 0; i < perf_loops; i++) {
230 for (j = 0; j < nnodes; j++)
231 insert_cached(nodes + j, &root);
232 for (j = 0; j < nnodes; j++)
233 erase_cached(nodes + j, &root);
234 }
235
236 time2 = get_cycles();
237 time = time2 - time1;
238
239 time = div_u64(time, perf_loops);
240 printk(" -> test 2 (latency of nnodes cached insert+delete): %llu cycles\n",
241 (unsigned long long)time);
242
243 for (i = 0; i < nnodes; i++)
244 insert(nodes + i, &root);
245
246 time1 = get_cycles();
247
248 for (i = 0; i < perf_loops; i++) {
249 for (node = rb_first(&root.rb_root); node; node = rb_next(node))
250 ;
251 }
252
253 time2 = get_cycles();
254 time = time2 - time1;
255
256 time = div_u64(time, perf_loops);
257 printk(" -> test 3 (latency of inorder traversal): %llu cycles\n",
258 (unsigned long long)time);
259
260 time1 = get_cycles();
261
262 for (i = 0; i < perf_loops; i++)
263 node = rb_first(&root.rb_root);
264
265 time2 = get_cycles();
266 time = time2 - time1;
267
268 time = div_u64(time, perf_loops);
269 printk(" -> test 4 (latency to fetch first node)\n");
270 printk(" non-cached: %llu cycles\n", (unsigned long long)time);
271
272 time1 = get_cycles();
273
274 for (i = 0; i < perf_loops; i++)
275 node = rb_first_cached(&root);
276
277 time2 = get_cycles();
278 time = time2 - time1;
279
280 time = div_u64(time, perf_loops);
281 printk(" cached: %llu cycles\n", (unsigned long long)time);
282
283 for (i = 0; i < nnodes; i++)
284 erase(nodes + i, &root);
285
286 /* run checks */
287 for (i = 0; i < check_loops; i++) {
288 init();
289 for (j = 0; j < nnodes; j++) {
290 check(j);
291 insert(nodes + j, &root);
292 }
293 for (j = 0; j < nnodes; j++) {
294 check(nnodes - j);
295 erase(nodes + j, &root);
296 }
297 check(0);
298 }
299
300 return 0;
301}
302
303static int augmented_check(void)
304{
305 int i, j;
306 cycles_t time1, time2, time;
307
308 printk(KERN_ALERT "augmented rbtree testing");
309
310 init();
311
312 time1 = get_cycles();
313
314 for (i = 0; i < perf_loops; i++) {
315 for (j = 0; j < nnodes; j++)
316 insert_augmented(nodes + j, &root);
317 for (j = 0; j < nnodes; j++)
318 erase_augmented(nodes + j, &root);
319 }
320
321 time2 = get_cycles();
322 time = time2 - time1;
323
324 time = div_u64(time, perf_loops);
325 printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n", (unsigned long long)time);
326
327 time1 = get_cycles();
328
329 for (i = 0; i < perf_loops; i++) {
330 for (j = 0; j < nnodes; j++)
331 insert_augmented_cached(nodes + j, &root);
332 for (j = 0; j < nnodes; j++)
333 erase_augmented_cached(nodes + j, &root);
334 }
335
336 time2 = get_cycles();
337 time = time2 - time1;
338
339 time = div_u64(time, perf_loops);
340 printk(" -> test 2 (latency of nnodes cached insert+delete): %llu cycles\n", (unsigned long long)time);
341
342 for (i = 0; i < check_loops; i++) {
343 init();
344 for (j = 0; j < nnodes; j++) {
345 check_augmented(j);
346 insert_augmented(nodes + j, &root);
347 }
348 for (j = 0; j < nnodes; j++) {
349 check_augmented(nnodes - j);
350 erase_augmented(nodes + j, &root);
351 }
352 check_augmented(0);
353 }
354
355 return 0;
356}
357
358static int __init rbtree_test_init(void)
359{
360 nodes = kmalloc_objs(*nodes, nnodes);
361 if (!nodes)
362 return -ENOMEM;
363
364 prandom_seed_state(&rnd, seed);
365
366 basic_check();
367 augmented_check();
368
369 kfree(nodes);
370
371 return -EAGAIN; /* Fail will directly unload the module */
372}
373
374static void __exit rbtree_test_exit(void)
375{
376 printk(KERN_ALERT "test exit\n");
377}
378
379module_init(rbtree_test_init)
380module_exit(rbtree_test_exit)
381
382MODULE_LICENSE("GPL");
383MODULE_AUTHOR("Michel Lespinasse");
384MODULE_DESCRIPTION("Red Black Tree test");