prio_tree.c 12 KB

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