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 /*
 308  * Following functions help to enumerate all prio_tree_nodes in the tree that
 309  * overlap with the input interval X [radix_index, heap_index]. The enumeration
 310  * takes O(log n + m) time where 'log n' is the height of the tree (which is
 311  * proportional to # of bits required to represent the maximum heap_index) and
 312  * 'm' is the number of prio_tree_nodes that overlap the interval X.
 313  */
 314
 315 static struct prio_tree_node *prio_tree_left(struct prio_tree_iter *iter,
 316                 unsigned long *r_index, unsigned long *h_index)
 317 {
 318         if (prio_tree_left_empty(iter->cur))
 319                 return NULL;
 320
 321         get_index(iter->root, iter->cur->left, r_index, h_index);
 322
 323         if (iter->r_index <= *h_index) {
 324                 iter->cur = iter->cur->left;
 325                 iter->mask >>= 1;
 326                 if (iter->mask) {
 327                         if (iter->size_level)
 328                                 iter->size_level++;
 329                 } else {
 330                         if (iter->size_level) {
 331                                 BUG_ON(!prio_tree_left_empty(iter->cur));
 332                                 BUG_ON(!prio_tree_right_empty(iter->cur));
 333                                 iter->size_level++;
 334                                 iter->mask = ULONG_MAX;
 335                         } else {
 336                                 iter->size_level = 1;
 337                                 iter->mask = 1UL << (BITS_PER_LONG - 1);
 338                         }
 339                 }
 340                 return iter->cur;
 341         }
 342
 343         return NULL;
 344 }
 345
 346 static struct prio_tree_node *prio_tree_right(struct prio_tree_iter *iter,
 347                 unsigned long *r_index, unsigned long *h_index)
 348 {
 349         unsigned long value;
 350
 351         if (prio_tree_right_empty(iter->cur))
 352                 return NULL;
 353
 354         if (iter->size_level)
 355                 value = iter->value;
 356         else
 357                 value = iter->value | iter->mask;
 358
 359         if (iter->h_index < value)
 360                 return NULL;
 361
 362         get_index(iter->root, iter->cur->right, r_index, h_index);
 363
 364         if (iter->r_index <= *h_index) {
 365                 iter->cur = iter->cur->right;
 366                 iter->mask >>= 1;
 367                 iter->value = value;
 368                 if (iter->mask) {
 369                         if (iter->size_level)
 370                                 iter->size_level++;
 371                 } else {
 372                         if (iter->size_level) {
 373                                 BUG_ON(!prio_tree_left_empty(iter->cur));
 374                                 BUG_ON(!prio_tree_right_empty(iter->cur));
 375                                 iter->size_level++;
 376                                 iter->mask = ULONG_MAX;
 377                         } else {
 378                                 iter->size_level = 1;
 379                                 iter->mask = 1UL << (BITS_PER_LONG - 1);
 380                         }
 381                 }
 382                 return iter->cur;
 383         }
 384
 385         return NULL;
 386 }
 387
 388 static struct prio_tree_node *prio_tree_parent(struct prio_tree_iter *iter)
 389 {
 390         iter->cur = iter->cur->parent;
 391         if (iter->mask == ULONG_MAX)
 392                 iter->mask = 1UL;
 393         else if (iter->size_level == 1)
 394                 iter->mask = 1UL;
 395         else
 396                 iter->mask <<= 1;
 397         if (iter->size_level)
 398                 iter->size_level--;
 399         if (!iter->size_level && (iter->value & iter->mask))
 400                 iter->value ^= iter->mask;
 401         return iter->cur;
 402 }
 403
 404 static inline int overlap(struct prio_tree_iter *iter,
 405                 unsigned long r_index, unsigned long h_index)
 406 {
 407         return iter->h_index >= r_index && iter->r_index <= h_index;
 408 }
 409
 410 /*
 411  * prio_tree_first:
 412  *
 413  * Get the first prio_tree_node that overlaps with the interval [radix_index,
 414  * heap_index]. Note that always radix_index <= heap_index. We do a pre-order
 415  * traversal of the tree.
 416  */
 417 static struct prio_tree_node *prio_tree_first(struct prio_tree_iter *iter)
 418 {
 419         struct prio_tree_root *root;
 420         unsigned long r_index, h_index;
 421
 422         INIT_PRIO_TREE_ITER(iter);
 423
 424         root = iter->root;
 425         if (prio_tree_empty(root))
 426                 return NULL;
 427
 428         get_index(root, root->prio_tree_node, &r_index, &h_index);
 429
 430         if (iter->r_index > h_index)
 431                 return NULL;
 432
 433         iter->mask = 1UL << (root->index_bits - 1);
 434         iter->cur = root->prio_tree_node;
 435
 436         while (1) {
 437                 if (overlap(iter, r_index, h_index))
 438                         return iter->cur;
 439
 440                 if (prio_tree_left(iter, &r_index, &h_index))
 441                         continue;
 442
 443                 if (prio_tree_right(iter, &r_index, &h_index))
 444                         continue;
 445
 446                 break;
 447         }
 448         return NULL;
 449 }
 450
 451 /*
 452  * prio_tree_next:
 453  *
 454  * Get the next prio_tree_node that overlaps with the input interval in iter
 455  */
 456 struct prio_tree_node *prio_tree_next(struct prio_tree_iter *iter)
 457 {
 458         unsigned long r_index, h_index;
 459
 460         if (iter->cur == NULL)
 461                 return prio_tree_first(iter);
 462
 463 repeat:
 464         while (prio_tree_left(iter, &r_index, &h_index))
 465                 if (overlap(iter, r_index, h_index))
 466                         return iter->cur;
 467
 468         while (!prio_tree_right(iter, &r_index, &h_index)) {
 469                 while (!prio_tree_root(iter->cur) &&
 470                                 iter->cur->parent->right == iter->cur)
 471                         prio_tree_parent(iter);
 472
 473                 if (prio_tree_root(iter->cur))
 474                         return NULL;
 475
 476                 prio_tree_parent(iter);
 477         }
 478
 479         if (overlap(iter, r_index, h_index))
 480                 return iter->cur;
 481
 482         goto repeat;
 483 }