888e8b3a97ea299e318f1cd79418269c95d696e9
[linux-2.6.git] / lib / prio_tree.c
1 /*
2  * lib/prio_tree.c - priority search tree
3  *
4  * Copyright (C) 2004, Rajesh Venkatasubramanian <vrajesh@umich.edu>
5  *
6  * This file is released under the GPL v2.
7  *
8  * Based on the radix priority search tree proposed by Edward M. McCreight
9  * SIAM Journal of Computing, vol. 14, no.2, pages 257-276, May 1985
10  *
11  * 02Feb2004    Initial version
12  */
13
14 #include <linux/init.h>
15 #include <linux/mm.h>
16 #include <linux/prio_tree.h>
17
18 /*
19  * A clever mix of heap and radix trees forms a radix priority search tree (PST)
20  * which is useful for storing intervals, e.g, we can consider a vma as a closed
21  * interval of file pages [offset_begin, offset_end], and store all vmas that
22  * map a file in a PST. Then, using the PST, we can answer a stabbing query,
23  * i.e., selecting a set of stored intervals (vmas) that overlap with (map) a
24  * given input interval X (a set of consecutive file pages), in "O(log n + m)"
25  * time where 'log n' is the height of the PST, and 'm' is the number of stored
26  * intervals (vmas) that overlap (map) with the input interval X (the set of
27  * consecutive file pages).
28  *
29  * In our implementation, we store closed intervals of the form [radix_index,
30  * heap_index]. We assume that always radix_index <= heap_index. McCreight's PST
31  * is designed for storing intervals with unique radix indices, i.e., each
32  * interval have different radix_index. However, this limitation can be easily
33  * overcome by using the size, i.e., heap_index - radix_index, as part of the
34  * index, so we index the tree using [(radix_index,size), heap_index].
35  *
36  * When the above-mentioned indexing scheme is used, theoretically, in a 32 bit
37  * machine, the maximum height of a PST can be 64. We can use a balanced version
38  * of the priority search tree to optimize the tree height, but the balanced
39  * tree proposed by McCreight is too complex and memory-hungry for our purpose.
40  */
41
42 /*
43  * The following macros are used for implementing prio_tree for i_mmap
44  */
45
46 #define RADIX_INDEX(vma)  ((vma)->vm_pgoff)
47 #define VMA_SIZE(vma)     (((vma)->vm_end - (vma)->vm_start) >> PAGE_SHIFT)
48 /* avoid overflow */
49 #define HEAP_INDEX(vma)   ((vma)->vm_pgoff + (VMA_SIZE(vma) - 1))
50
51
52 static void get_index(const struct prio_tree_root *root,
53     const struct prio_tree_node *node,
54     unsigned long *radix, unsigned long *heap)
55 {
56         if (root->raw) {
57                 struct vm_area_struct *vma = prio_tree_entry(
58                     node, struct vm_area_struct, shared.prio_tree_node);
59
60                 *radix = RADIX_INDEX(vma);
61                 *heap = HEAP_INDEX(vma);
62         }
63         else {
64                 *radix = node->start;
65                 *heap = node->last;
66         }
67 }
68
69 static unsigned long index_bits_to_maxindex[BITS_PER_LONG];
70
71 void __init prio_tree_init(void)
72 {
73         unsigned int i;
74
75         for (i = 0; i < ARRAY_SIZE(index_bits_to_maxindex) - 1; i++)
76                 index_bits_to_maxindex[i] = (1UL << (i + 1)) - 1;
77         index_bits_to_maxindex[ARRAY_SIZE(index_bits_to_maxindex) - 1] = ~0UL;
78 }
79
80 /*
81  * Maximum heap_index that can be stored in a PST with index_bits bits
82  */
83 static inline unsigned long prio_tree_maxindex(unsigned int bits)
84 {
85         return index_bits_to_maxindex[bits - 1];
86 }
87
88 /*
89  * Extend a priority search tree so that it can store a node with heap_index
90  * max_heap_index. In the worst case, this algorithm takes O((log n)^2).
91  * However, this function is used rarely and the common case performance is
92  * not bad.
93  */
94 static struct prio_tree_node *prio_tree_expand(struct prio_tree_root *root,
95                 struct prio_tree_node *node, unsigned long max_heap_index)
96 {
97         struct prio_tree_node *first = NULL, *prev, *last = NULL;
98
99         if (max_heap_index > prio_tree_maxindex(root->index_bits))
100                 root->index_bits++;
101
102         while (max_heap_index > prio_tree_maxindex(root->index_bits)) {
103                 root->index_bits++;
104
105                 if (prio_tree_empty(root))
106                         continue;
107
108                 if (first == NULL) {
109                         first = root->prio_tree_node;
110                         prio_tree_remove(root, root->prio_tree_node);
111                         INIT_PRIO_TREE_NODE(first);
112                         last = first;
113                 } else {
114                         prev = last;
115                         last = root->prio_tree_node;
116                         prio_tree_remove(root, root->prio_tree_node);
117                         INIT_PRIO_TREE_NODE(last);
118                         prev->left = last;
119                         last->parent = prev;
120                 }
121         }
122
123         INIT_PRIO_TREE_NODE(node);
124
125         if (first) {
126                 node->left = first;
127                 first->parent = node;
128         } else
129                 last = node;
130
131         if (!prio_tree_empty(root)) {
132                 last->left = root->prio_tree_node;
133                 last->left->parent = last;
134         }
135
136         root->prio_tree_node = node;
137         return node;
138 }
139
140 /*
141  * Replace a prio_tree_node with a new node and return the old node
142  */
143 struct prio_tree_node *prio_tree_replace(struct prio_tree_root *root,
144                 struct prio_tree_node *old, struct prio_tree_node *node)
145 {
146         INIT_PRIO_TREE_NODE(node);
147
148         if (prio_tree_root(old)) {
149                 BUG_ON(root->prio_tree_node != old);
150                 /*
151                  * We can reduce root->index_bits here. However, it is complex
152                  * and does not help much to improve performance (IMO).
153                  */
154                 root->prio_tree_node = node;
155         } else {
156                 node->parent = old->parent;
157                 if (old->parent->left == old)
158                         old->parent->left = node;
159                 else
160                         old->parent->right = node;
161         }
162
163         if (!prio_tree_left_empty(old)) {
164                 node->left = old->left;
165                 old->left->parent = node;
166         }
167
168         if (!prio_tree_right_empty(old)) {
169                 node->right = old->right;
170                 old->right->parent = node;
171         }
172
173         return old;
174 }
175
176 /*
177  * Insert a prio_tree_node @node into a radix priority search tree @root. The
178  * algorithm typically takes O(log n) time where 'log n' is the number of bits
179  * required to represent the maximum heap_index. In the worst case, the algo
180  * can take O((log n)^2) - check prio_tree_expand.
181  *
182  * If a prior node with same radix_index and heap_index is already found in
183  * the tree, then returns the address of the prior node. Otherwise, inserts
184  * @node into the tree and returns @node.
185  */
186 struct prio_tree_node *prio_tree_insert(struct prio_tree_root *root,
187                 struct prio_tree_node *node)
188 {
189         struct prio_tree_node *cur, *res = node;
190         unsigned long radix_index, heap_index;
191         unsigned long r_index, h_index, index, mask;
192         int size_flag = 0;
193
194         get_index(root, node, &radix_index, &heap_index);
195
196         if (prio_tree_empty(root) ||
197                         heap_index > prio_tree_maxindex(root->index_bits))
198                 return prio_tree_expand(root, node, heap_index);
199
200         cur = root->prio_tree_node;
201         mask = 1UL << (root->index_bits - 1);
202
203         while (mask) {
204                 get_index(root, cur, &r_index, &h_index);
205
206                 if (r_index == radix_index && h_index == heap_index)
207                         return cur;
208
209                 if (h_index < heap_index ||
210                     (h_index == heap_index && r_index > radix_index)) {
211                         struct prio_tree_node *tmp = node;
212                         node = prio_tree_replace(root, cur, node);
213                         cur = tmp;
214                         /* swap indices */
215                         index = r_index;
216                         r_index = radix_index;
217                         radix_index = index;
218                         index = h_index;
219                         h_index = heap_index;
220                         heap_index = index;
221                 }
222
223                 if (size_flag)
224                         index = heap_index - radix_index;
225                 else
226                         index = radix_index;
227
228                 if (index & mask) {
229                         if (prio_tree_right_empty(cur)) {
230                                 INIT_PRIO_TREE_NODE(node);
231                                 cur->right = node;
232                                 node->parent = cur;
233                                 return res;
234                         } else
235                                 cur = cur->right;
236                 } else {
237                         if (prio_tree_left_empty(cur)) {
238                                 INIT_PRIO_TREE_NODE(node);
239                                 cur->left = node;
240                                 node->parent = cur;
241                                 return res;
242                         } else
243                                 cur = cur->left;
244                 }
245
246                 mask >>= 1;
247
248                 if (!mask) {
249                         mask = 1UL << (BITS_PER_LONG - 1);
250                         size_flag = 1;
251                 }
252         }
253         /* Should not reach here */
254         BUG();
255         return NULL;
256 }
257
258 /*
259  * Remove a prio_tree_node @node from a radix priority search tree @root. The
260  * algorithm takes O(log n) time where 'log n' is the number of bits required
261  * to represent the maximum heap_index.
262  */
263 void prio_tree_remove(struct prio_tree_root *root, struct prio_tree_node *node)
264 {
265         struct prio_tree_node *cur;
266         unsigned long r_index, h_index_right, h_index_left;
267
268         cur = node;
269
270         while (!prio_tree_left_empty(cur) || !prio_tree_right_empty(cur)) {
271                 if (!prio_tree_left_empty(cur))
272                         get_index(root, cur->left, &r_index, &h_index_left);
273                 else {
274                         cur = cur->right;
275                         continue;
276                 }
277
278                 if (!prio_tree_right_empty(cur))
279                         get_index(root, cur->right, &r_index, &h_index_right);
280                 else {
281                         cur = cur->left;
282                         continue;
283                 }
284
285                 /* both h_index_left and h_index_right cannot be 0 */
286                 if (h_index_left >= h_index_right)
287                         cur = cur->left;
288                 else
289                         cur = cur->right;
290         }
291
292         if (prio_tree_root(cur)) {
293                 BUG_ON(root->prio_tree_node != cur);
294                 __INIT_PRIO_TREE_ROOT(root, root->raw);
295                 return;
296         }
297
298         if (cur->parent->right == cur)
299                 cur->parent->right = cur->parent;
300         else
301                 cur->parent->left = cur->parent;
302
303         while (cur != node)
304                 cur = prio_tree_replace(root, cur->parent, cur);
305 }
306
307 static void iter_walk_down(struct prio_tree_iter *iter)
308 {
309         iter->mask >>= 1;
310         if (iter->mask) {
311                 if (iter->size_level)
312                         iter->size_level++;
313                 return;
314         }
315
316         if (iter->size_level) {
317                 BUG_ON(!prio_tree_left_empty(iter->cur));
318                 BUG_ON(!prio_tree_right_empty(iter->cur));
319                 iter->size_level++;
320                 iter->mask = ULONG_MAX;
321         } else {
322                 iter->size_level = 1;
323                 iter->mask = 1UL << (BITS_PER_LONG - 1);
324         }
325 }
326
327 static void iter_walk_up(struct prio_tree_iter *iter)
328 {
329         if (iter->mask == ULONG_MAX)
330                 iter->mask = 1UL;
331         else if (iter->size_level == 1)
332                 iter->mask = 1UL;
333         else
334                 iter->mask <<= 1;
335         if (iter->size_level)
336                 iter->size_level--;
337         if (!iter->size_level && (iter->value & iter->mask))
338                 iter->value ^= iter->mask;
339 }
340
341 /*
342  * Following functions help to enumerate all prio_tree_nodes in the tree that
343  * overlap with the input interval X [radix_index, heap_index]. The enumeration
344  * takes O(log n + m) time where 'log n' is the height of the tree (which is
345  * proportional to # of bits required to represent the maximum heap_index) and
346  * 'm' is the number of prio_tree_nodes that overlap the interval X.
347  */
348
349 static struct prio_tree_node *prio_tree_left(struct prio_tree_iter *iter,
350                 unsigned long *r_index, unsigned long *h_index)
351 {
352         if (prio_tree_left_empty(iter->cur))
353                 return NULL;
354
355         get_index(iter->root, iter->cur->left, r_index, h_index);
356
357         if (iter->r_index <= *h_index) {
358                 iter->cur = iter->cur->left;
359                 iter_walk_down(iter);
360                 return iter->cur;
361         }
362
363         return NULL;
364 }
365
366 static struct prio_tree_node *prio_tree_right(struct prio_tree_iter *iter,
367                 unsigned long *r_index, unsigned long *h_index)
368 {
369         unsigned long value;
370
371         if (prio_tree_right_empty(iter->cur))
372                 return NULL;
373
374         if (iter->size_level)
375                 value = iter->value;
376         else
377                 value = iter->value | iter->mask;
378
379         if (iter->h_index < value)
380                 return NULL;
381
382         get_index(iter->root, iter->cur->right, r_index, h_index);
383
384         if (iter->r_index <= *h_index) {
385                 iter->cur = iter->cur->right;
386                 iter_walk_down(iter);
387                 return iter->cur;
388         }
389
390         return NULL;
391 }
392
393 static struct prio_tree_node *prio_tree_parent(struct prio_tree_iter *iter)
394 {
395         iter->cur = iter->cur->parent;
396         iter_walk_up(iter);
397         return iter->cur;
398 }
399
400 static inline int overlap(struct prio_tree_iter *iter,
401                 unsigned long r_index, unsigned long h_index)
402 {
403         return iter->h_index >= r_index && iter->r_index <= h_index;
404 }
405
406 /*
407  * prio_tree_first:
408  *
409  * Get the first prio_tree_node that overlaps with the interval [radix_index,
410  * heap_index]. Note that always radix_index <= heap_index. We do a pre-order
411  * traversal of the tree.
412  */
413 static struct prio_tree_node *prio_tree_first(struct prio_tree_iter *iter)
414 {
415         struct prio_tree_root *root;
416         unsigned long r_index, h_index;
417
418         INIT_PRIO_TREE_ITER(iter);
419
420         root = iter->root;
421         if (prio_tree_empty(root))
422                 return NULL;
423
424         get_index(root, root->prio_tree_node, &r_index, &h_index);
425
426         if (iter->r_index > h_index)
427                 return NULL;
428
429         iter->mask = 1UL << (root->index_bits - 1);
430         iter->cur = root->prio_tree_node;
431
432         while (1) {
433                 if (overlap(iter, r_index, h_index))
434                         return iter->cur;
435
436                 if (prio_tree_left(iter, &r_index, &h_index))
437                         continue;
438
439                 if (prio_tree_right(iter, &r_index, &h_index))
440                         continue;
441
442                 break;
443         }
444         return NULL;
445 }
446
447 /*
448  * prio_tree_next:
449  *
450  * Get the next prio_tree_node that overlaps with the input interval in iter
451  */
452 struct prio_tree_node *prio_tree_next(struct prio_tree_iter *iter)
453 {
454         unsigned long r_index, h_index;
455
456         if (iter->cur == NULL)
457                 return prio_tree_first(iter);
458
459 repeat:
460         while (prio_tree_left(iter, &r_index, &h_index))
461                 if (overlap(iter, r_index, h_index))
462                         return iter->cur;
463
464         while (!prio_tree_right(iter, &r_index, &h_index)) {
465                 while (!prio_tree_root(iter->cur) &&
466                                 iter->cur->parent->right == iter->cur)
467                         prio_tree_parent(iter);
468
469                 if (prio_tree_root(iter->cur))
470                         return NULL;
471
472                 prio_tree_parent(iter);
473         }
474
475         if (overlap(iter, r_index, h_index))
476                 return iter->cur;
477
478         goto repeat;
479 }