Merge branch 'for-next' of git://git.kernel.org/pub/scm/linux/kernel/git/cooloney...
[platform/adaptation/renesas_rcar/renesas_kernel.git] / lib / rbtree_test.c
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
10 struct test_node {
11         struct rb_node rb;
12         u32 key;
13
14         /* following fields used for testing augmented rbtree functionality */
15         u32 val;
16         u32 augmented;
17 };
18
19 static struct rb_root root = RB_ROOT;
20 static struct test_node nodes[NODES];
21
22 static struct rnd_state rnd;
23
24 static 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
41 static inline void erase(struct test_node *node, struct rb_root *root)
42 {
43         rb_erase(&node->rb, root);
44 }
45
46 static 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
64 RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
65                      u32, augmented, augment_recompute)
66
67 static 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
90 static void erase_augmented(struct test_node *node, struct rb_root *root)
91 {
92         rb_erase_augmented(&node->rb, root, &augment_callbacks);
93 }
94
95 static 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
104 static bool is_red(struct rb_node *rb)
105 {
106         return !(rb->__rb_parent_color & 1);
107 }
108
109 static 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
117 static void check_postorder(int nr_nodes)
118 {
119         struct rb_node *rb;
120         int count = 0;
121         for (rb = rb_first_postorder(&root); rb; rb = rb_next_postorder(rb))
122                 count++;
123
124         WARN_ON_ONCE(count != nr_nodes);
125 }
126
127 static void check(int nr_nodes)
128 {
129         struct rb_node *rb;
130         int count = 0, blacks = 0;
131         u32 prev_key = 0;
132
133         for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
134                 struct test_node *node = rb_entry(rb, struct test_node, rb);
135                 WARN_ON_ONCE(node->key < prev_key);
136                 WARN_ON_ONCE(is_red(rb) &&
137                              (!rb_parent(rb) || is_red(rb_parent(rb))));
138                 if (!count)
139                         blacks = black_path_count(rb);
140                 else
141                         WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
142                                      blacks != black_path_count(rb));
143                 prev_key = node->key;
144                 count++;
145         }
146
147         WARN_ON_ONCE(count != nr_nodes);
148         WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);
149
150         check_postorder(nr_nodes);
151 }
152
153 static void check_augmented(int nr_nodes)
154 {
155         struct rb_node *rb;
156
157         check(nr_nodes);
158         for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
159                 struct test_node *node = rb_entry(rb, struct test_node, rb);
160                 WARN_ON_ONCE(node->augmented != augment_recompute(node));
161         }
162 }
163
164 static int __init rbtree_test_init(void)
165 {
166         int i, j;
167         cycles_t time1, time2, time;
168
169         printk(KERN_ALERT "rbtree testing");
170
171         prandom_seed_state(&rnd, 3141592653589793238ULL);
172         init();
173
174         time1 = get_cycles();
175
176         for (i = 0; i < PERF_LOOPS; i++) {
177                 for (j = 0; j < NODES; j++)
178                         insert(nodes + j, &root);
179                 for (j = 0; j < NODES; j++)
180                         erase(nodes + j, &root);
181         }
182
183         time2 = get_cycles();
184         time = time2 - time1;
185
186         time = div_u64(time, PERF_LOOPS);
187         printk(" -> %llu cycles\n", (unsigned long long)time);
188
189         for (i = 0; i < CHECK_LOOPS; i++) {
190                 init();
191                 for (j = 0; j < NODES; j++) {
192                         check(j);
193                         insert(nodes + j, &root);
194                 }
195                 for (j = 0; j < NODES; j++) {
196                         check(NODES - j);
197                         erase(nodes + j, &root);
198                 }
199                 check(0);
200         }
201
202         printk(KERN_ALERT "augmented rbtree testing");
203
204         init();
205
206         time1 = get_cycles();
207
208         for (i = 0; i < PERF_LOOPS; i++) {
209                 for (j = 0; j < NODES; j++)
210                         insert_augmented(nodes + j, &root);
211                 for (j = 0; j < NODES; j++)
212                         erase_augmented(nodes + j, &root);
213         }
214
215         time2 = get_cycles();
216         time = time2 - time1;
217
218         time = div_u64(time, PERF_LOOPS);
219         printk(" -> %llu cycles\n", (unsigned long long)time);
220
221         for (i = 0; i < CHECK_LOOPS; i++) {
222                 init();
223                 for (j = 0; j < NODES; j++) {
224                         check_augmented(j);
225                         insert_augmented(nodes + j, &root);
226                 }
227                 for (j = 0; j < NODES; j++) {
228                         check_augmented(NODES - j);
229                         erase_augmented(nodes + j, &root);
230                 }
231                 check_augmented(0);
232         }
233
234         return -EAGAIN; /* Fail will directly unload the module */
235 }
236
237 static void __exit rbtree_test_exit(void)
238 {
239         printk(KERN_ALERT "test exit\n");
240 }
241
242 module_init(rbtree_test_init)
243 module_exit(rbtree_test_exit)
244
245 MODULE_LICENSE("GPL");
246 MODULE_AUTHOR("Michel Lespinasse");
247 MODULE_DESCRIPTION("Red Black Tree test");