xref: /linux/lib/interval_tree_test.c (revision 0000d9ccbcfa90411c88f70850501723389312b9)
1 // SPDX-License-Identifier: GPL-2.0-only
2 #include <linux/module.h>
3 #include <linux/moduleparam.h>
4 #include <linux/interval_tree.h>
5 #include <linux/prandom.h>
6 #include <linux/slab.h>
7 #include <linux/printk.h>
8 #include <asm/timex.h>
9 #include <linux/bitmap.h>
10 #include <linux/maple_tree.h>
11 
12 #define __param(type, name, init, msg)		\
13 	static type name = init;		\
14 	module_param(name, type, 0444);		\
15 	MODULE_PARM_DESC(name, msg);
16 
17 __param(int, nnodes, 100, "Number of nodes in the interval tree");
18 __param(int, perf_loops, 1000, "Number of iterations modifying the tree");
19 
20 __param(int, nsearches, 100, "Number of searches to the interval tree");
21 __param(int, search_loops, 1000, "Number of iterations searching the tree");
22 __param(bool, search_all, false, "Searches will iterate all nodes in the tree");
23 
24 __param(uint, max_endpoint, ~0, "Largest value for the interval's endpoint");
25 __param(ullong, seed, 3141592653589793238ULL, "Random seed");
26 
27 static struct rb_root_cached root = RB_ROOT_CACHED;
28 static struct interval_tree_node *nodes = NULL;
29 static u32 *queries = NULL;
30 
31 static struct rnd_state rnd;
32 
33 static inline unsigned long
34 search(struct rb_root_cached *root, unsigned long start, unsigned long last)
35 {
36 	struct interval_tree_node *node;
37 	unsigned long results = 0;
38 
39 	for (node = interval_tree_iter_first(root, start, last); node;
40 	     node = interval_tree_iter_next(node, start, last))
41 		results++;
42 	return results;
43 }
44 
45 static void init(void)
46 {
47 	int i;
48 
49 	for (i = 0; i < nnodes; i++) {
50 		u32 b = (prandom_u32_state(&rnd) >> 4) % max_endpoint;
51 		u32 a = (prandom_u32_state(&rnd) >> 4) % b;
52 
53 		nodes[i].start = a;
54 		nodes[i].last = b;
55 	}
56 
57 	/*
58 	 * Limit the search scope to what the user defined.
59 	 * Otherwise we are merely measuring empty walks,
60 	 * which is pointless.
61 	 */
62 	for (i = 0; i < nsearches; i++)
63 		queries[i] = (prandom_u32_state(&rnd) >> 4) % max_endpoint;
64 }
65 
66 static int basic_check(void)
67 {
68 	int i, j;
69 	cycles_t time1, time2, time;
70 
71 	printk(KERN_ALERT "interval tree insert/remove");
72 
73 	init();
74 
75 	time1 = get_cycles();
76 
77 	for (i = 0; i < perf_loops; i++) {
78 		for (j = 0; j < nnodes; j++)
79 			interval_tree_insert(nodes + j, &root);
80 		for (j = 0; j < nnodes; j++)
81 			interval_tree_remove(nodes + j, &root);
82 	}
83 
84 	time2 = get_cycles();
85 	time = time2 - time1;
86 
87 	time = div_u64(time, perf_loops);
88 	printk(" -> %llu cycles\n", (unsigned long long)time);
89 
90 	return 0;
91 }
92 
93 static int search_check(void)
94 {
95 	int i, j;
96 	unsigned long results;
97 	cycles_t time1, time2, time;
98 
99 	printk(KERN_ALERT "interval tree search");
100 
101 	init();
102 
103 	for (j = 0; j < nnodes; j++)
104 		interval_tree_insert(nodes + j, &root);
105 
106 	time1 = get_cycles();
107 
108 	results = 0;
109 	for (i = 0; i < search_loops; i++)
110 		for (j = 0; j < nsearches; j++) {
111 			unsigned long start = search_all ? 0 : queries[j];
112 			unsigned long last = search_all ? max_endpoint : queries[j];
113 
114 			results += search(&root, start, last);
115 		}
116 
117 	time2 = get_cycles();
118 	time = time2 - time1;
119 
120 	time = div_u64(time, search_loops);
121 	results = div_u64(results, search_loops);
122 	printk(" -> %llu cycles (%lu results)\n",
123 	       (unsigned long long)time, results);
124 
125 	for (j = 0; j < nnodes; j++)
126 		interval_tree_remove(nodes + j, &root);
127 
128 	return 0;
129 }
130 
131 static int intersection_range_check(void)
132 {
133 	int i, j, k;
134 	unsigned long start, last;
135 	struct interval_tree_node *node;
136 	unsigned long *intxn1;
137 	unsigned long *intxn2;
138 
139 	printk(KERN_ALERT "interval tree iteration\n");
140 
141 	intxn1 = bitmap_alloc(nnodes, GFP_KERNEL);
142 	if (!intxn1) {
143 		WARN_ON_ONCE("Failed to allocate intxn1\n");
144 		return -ENOMEM;
145 	}
146 
147 	intxn2 = bitmap_alloc(nnodes, GFP_KERNEL);
148 	if (!intxn2) {
149 		WARN_ON_ONCE("Failed to allocate intxn2\n");
150 		bitmap_free(intxn1);
151 		return -ENOMEM;
152 	}
153 
154 	for (i = 0; i < search_loops; i++) {
155 		/* Initialize interval tree for each round */
156 		init();
157 		for (j = 0; j < nnodes; j++)
158 			interval_tree_insert(nodes + j, &root);
159 
160 		/* Let's try nsearches different ranges */
161 		for (k = 0; k < nsearches; k++) {
162 			/* Try whole range once */
163 			if (!k) {
164 				start = 0UL;
165 				last = ULONG_MAX;
166 			} else {
167 				last = (prandom_u32_state(&rnd) >> 4) % max_endpoint;
168 				start = (prandom_u32_state(&rnd) >> 4) % last;
169 			}
170 
171 			/* Walk nodes to mark intersection nodes */
172 			bitmap_zero(intxn1, nnodes);
173 			for (j = 0; j < nnodes; j++) {
174 				node = nodes + j;
175 
176 				if (start <= node->last && last >= node->start)
177 					bitmap_set(intxn1, j, 1);
178 			}
179 
180 			/* Iterate tree to clear intersection nodes */
181 			bitmap_zero(intxn2, nnodes);
182 			for (node = interval_tree_iter_first(&root, start, last); node;
183 			     node = interval_tree_iter_next(node, start, last))
184 				bitmap_set(intxn2, node - nodes, 1);
185 
186 			WARN_ON_ONCE(!bitmap_equal(intxn1, intxn2, nnodes));
187 		}
188 
189 		for (j = 0; j < nnodes; j++)
190 			interval_tree_remove(nodes + j, &root);
191 	}
192 
193 	bitmap_free(intxn1);
194 	bitmap_free(intxn2);
195 	return 0;
196 }
197 
198 #ifdef CONFIG_INTERVAL_TREE_SPAN_ITER
199 /*
200  * Helper function to get span of current position from maple tree point of
201  * view.
202  */
203 static void mas_cur_span(struct ma_state *mas, struct interval_tree_span_iter *state)
204 {
205 	unsigned long cur_start;
206 	unsigned long cur_last;
207 	int is_hole;
208 
209 	if (mas->status == ma_overflow)
210 		return;
211 
212 	/* walk to current position */
213 	state->is_hole = mas_walk(mas) ? 0 : 1;
214 
215 	cur_start = mas->index < state->first_index ?
216 			state->first_index : mas->index;
217 
218 	/* whether we have followers */
219 	do {
220 
221 		cur_last = mas->last > state->last_index ?
222 				state->last_index : mas->last;
223 
224 		is_hole = mas_next_range(mas, state->last_index) ? 0 : 1;
225 
226 	} while (mas->status != ma_overflow && is_hole == state->is_hole);
227 
228 	if (state->is_hole) {
229 		state->start_hole = cur_start;
230 		state->last_hole = cur_last;
231 	} else {
232 		state->start_used = cur_start;
233 		state->last_used = cur_last;
234 	}
235 
236 	/* advance position for next round */
237 	if (mas->status != ma_overflow)
238 		mas_set(mas, cur_last + 1);
239 }
240 
241 static int span_iteration_check(void)
242 {
243 	int i, j, k;
244 	unsigned long start, last;
245 	struct interval_tree_span_iter span, mas_span;
246 
247 	DEFINE_MTREE(tree);
248 
249 	MA_STATE(mas, &tree, 0, 0);
250 
251 	printk(KERN_ALERT "interval tree span iteration\n");
252 
253 	for (i = 0; i < search_loops; i++) {
254 		/* Initialize interval tree for each round */
255 		init();
256 		for (j = 0; j < nnodes; j++)
257 			interval_tree_insert(nodes + j, &root);
258 
259 		/* Put all the range into maple tree */
260 		mt_init_flags(&tree, MT_FLAGS_ALLOC_RANGE);
261 		mt_set_in_rcu(&tree);
262 
263 		for (j = 0; j < nnodes; j++)
264 			WARN_ON_ONCE(mtree_store_range(&tree, nodes[j].start,
265 					nodes[j].last, nodes + j, GFP_KERNEL));
266 
267 		/* Let's try nsearches different ranges */
268 		for (k = 0; k < nsearches; k++) {
269 			/* Try whole range once */
270 			if (!k) {
271 				start = 0UL;
272 				last = ULONG_MAX;
273 			} else {
274 				last = (prandom_u32_state(&rnd) >> 4) % max_endpoint;
275 				start = (prandom_u32_state(&rnd) >> 4) % last;
276 			}
277 
278 			mas_span.first_index = start;
279 			mas_span.last_index = last;
280 			mas_span.is_hole = -1;
281 			mas_set(&mas, start);
282 
283 			interval_tree_for_each_span(&span, &root, start, last) {
284 				mas_cur_span(&mas, &mas_span);
285 
286 				WARN_ON_ONCE(span.is_hole != mas_span.is_hole);
287 
288 				if (span.is_hole) {
289 					WARN_ON_ONCE(span.start_hole != mas_span.start_hole);
290 					WARN_ON_ONCE(span.last_hole != mas_span.last_hole);
291 				} else {
292 					WARN_ON_ONCE(span.start_used != mas_span.start_used);
293 					WARN_ON_ONCE(span.last_used != mas_span.last_used);
294 				}
295 			}
296 
297 		}
298 
299 		WARN_ON_ONCE(mas.status != ma_overflow);
300 
301 		/* Cleanup maple tree for each round */
302 		mtree_destroy(&tree);
303 		/* Cleanup interval tree for each round */
304 		for (j = 0; j < nnodes; j++)
305 			interval_tree_remove(nodes + j, &root);
306 	}
307 	return 0;
308 }
309 #else
310 static inline int span_iteration_check(void) {return 0; }
311 #endif
312 
313 static int interval_tree_test_init(void)
314 {
315 	if (nnodes <= 0) {
316 		pr_warn("nnodes must be positive\n");
317 		return -EINVAL;
318 	}
319 	if (nsearches <= 0) {
320 		pr_warn("nsearches must be positive\n");
321 		return -EINVAL;
322 	}
323 	if (perf_loops <= 0) {
324 		pr_warn("perf_loops must be positive\n");
325 		return -EINVAL;
326 	}
327 	if (search_loops <= 0) {
328 		pr_warn("search_loops must be positive\n");
329 		return -EINVAL;
330 	}
331 	if (max_endpoint < 2) {
332 		pr_warn("max_endpoint must be at least 2\n");
333 		return -EINVAL;
334 	}
335 
336 	nodes = kmalloc_objs(struct interval_tree_node, nnodes);
337 	if (!nodes)
338 		return -ENOMEM;
339 
340 	queries = kmalloc_array(nsearches, sizeof(int), GFP_KERNEL);
341 	if (!queries) {
342 		kfree(nodes);
343 		return -ENOMEM;
344 	}
345 
346 	prandom_seed_state(&rnd, seed);
347 
348 	basic_check();
349 	search_check();
350 	intersection_range_check();
351 	span_iteration_check();
352 
353 	kfree(queries);
354 	kfree(nodes);
355 
356 	return -EAGAIN; /* Fail will directly unload the module */
357 }
358 
359 static void interval_tree_test_exit(void)
360 {
361 	printk(KERN_ALERT "test exit\n");
362 }
363 
364 module_init(interval_tree_test_init)
365 module_exit(interval_tree_test_exit)
366 
367 MODULE_LICENSE("GPL");
368 MODULE_AUTHOR("Michel Lespinasse");
369 MODULE_DESCRIPTION("Interval Tree test");
370