interval-tree.c 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882
  1. /* SPDX-License-Identifier: GPL-2.0-or-later */
  2. #include "qemu/osdep.h"
  3. #include "qemu/interval-tree.h"
  4. #include "qemu/atomic.h"
  5. /*
  6. * Red Black Trees.
  7. *
  8. * For now, don't expose Linux Red-Black Trees separately, but retain the
  9. * separate type definitions to keep the implementation sane, and allow
  10. * the possibility of separating them later.
  11. *
  12. * Derived from include/linux/rbtree_augmented.h and its dependencies.
  13. */
  14. /*
  15. * red-black trees properties: https://en.wikipedia.org/wiki/Rbtree
  16. *
  17. * 1) A node is either red or black
  18. * 2) The root is black
  19. * 3) All leaves (NULL) are black
  20. * 4) Both children of every red node are black
  21. * 5) Every simple path from root to leaves contains the same number
  22. * of black nodes.
  23. *
  24. * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
  25. * consecutive red nodes in a path and every red node is therefore followed by
  26. * a black. So if B is the number of black nodes on every simple path (as per
  27. * 5), then the longest possible path due to 4 is 2B.
  28. *
  29. * We shall indicate color with case, where black nodes are uppercase and red
  30. * nodes will be lowercase. Unknown color nodes shall be drawn as red within
  31. * parentheses and have some accompanying text comment.
  32. *
  33. * Notes on lockless lookups:
  34. *
  35. * All stores to the tree structure (rb_left and rb_right) must be done using
  36. * WRITE_ONCE [qatomic_set for QEMU]. And we must not inadvertently cause
  37. * (temporary) loops in the tree structure as seen in program order.
  38. *
  39. * These two requirements will allow lockless iteration of the tree -- not
  40. * correct iteration mind you, tree rotations are not atomic so a lookup might
  41. * miss entire subtrees.
  42. *
  43. * But they do guarantee that any such traversal will only see valid elements
  44. * and that it will indeed complete -- does not get stuck in a loop.
  45. *
  46. * It also guarantees that if the lookup returns an element it is the 'correct'
  47. * one. But not returning an element does _NOT_ mean it's not present.
  48. *
  49. * NOTE:
  50. *
  51. * Stores to __rb_parent_color are not important for simple lookups so those
  52. * are left undone as of now. Nor did I check for loops involving parent
  53. * pointers.
  54. */
  55. typedef enum RBColor
  56. {
  57. RB_RED,
  58. RB_BLACK,
  59. } RBColor;
  60. typedef struct RBAugmentCallbacks {
  61. void (*propagate)(RBNode *node, RBNode *stop);
  62. void (*copy)(RBNode *old, RBNode *new);
  63. void (*rotate)(RBNode *old, RBNode *new);
  64. } RBAugmentCallbacks;
  65. static inline RBNode *rb_parent(const RBNode *n)
  66. {
  67. return (RBNode *)(n->rb_parent_color & ~1);
  68. }
  69. static inline RBNode *rb_red_parent(const RBNode *n)
  70. {
  71. return (RBNode *)n->rb_parent_color;
  72. }
  73. static inline RBColor pc_color(uintptr_t pc)
  74. {
  75. return (RBColor)(pc & 1);
  76. }
  77. static inline bool pc_is_red(uintptr_t pc)
  78. {
  79. return pc_color(pc) == RB_RED;
  80. }
  81. static inline bool pc_is_black(uintptr_t pc)
  82. {
  83. return !pc_is_red(pc);
  84. }
  85. static inline RBColor rb_color(const RBNode *n)
  86. {
  87. return pc_color(n->rb_parent_color);
  88. }
  89. static inline bool rb_is_red(const RBNode *n)
  90. {
  91. return pc_is_red(n->rb_parent_color);
  92. }
  93. static inline bool rb_is_black(const RBNode *n)
  94. {
  95. return pc_is_black(n->rb_parent_color);
  96. }
  97. static inline void rb_set_black(RBNode *n)
  98. {
  99. n->rb_parent_color |= RB_BLACK;
  100. }
  101. static inline void rb_set_parent_color(RBNode *n, RBNode *p, RBColor color)
  102. {
  103. n->rb_parent_color = (uintptr_t)p | color;
  104. }
  105. static inline void rb_set_parent(RBNode *n, RBNode *p)
  106. {
  107. rb_set_parent_color(n, p, rb_color(n));
  108. }
  109. static inline void rb_link_node(RBNode *node, RBNode *parent, RBNode **rb_link)
  110. {
  111. node->rb_parent_color = (uintptr_t)parent;
  112. node->rb_left = node->rb_right = NULL;
  113. qatomic_set(rb_link, node);
  114. }
  115. static RBNode *rb_next(RBNode *node)
  116. {
  117. RBNode *parent;
  118. /* OMIT: if empty node, return null. */
  119. /*
  120. * If we have a right-hand child, go down and then left as far as we can.
  121. */
  122. if (node->rb_right) {
  123. node = node->rb_right;
  124. while (node->rb_left) {
  125. node = node->rb_left;
  126. }
  127. return node;
  128. }
  129. /*
  130. * No right-hand children. Everything down and left is smaller than us,
  131. * so any 'next' node must be in the general direction of our parent.
  132. * Go up the tree; any time the ancestor is a right-hand child of its
  133. * parent, keep going up. First time it's a left-hand child of its
  134. * parent, said parent is our 'next' node.
  135. */
  136. while ((parent = rb_parent(node)) && node == parent->rb_right) {
  137. node = parent;
  138. }
  139. return parent;
  140. }
  141. static inline void rb_change_child(RBNode *old, RBNode *new,
  142. RBNode *parent, RBRoot *root)
  143. {
  144. if (!parent) {
  145. qatomic_set(&root->rb_node, new);
  146. } else if (parent->rb_left == old) {
  147. qatomic_set(&parent->rb_left, new);
  148. } else {
  149. qatomic_set(&parent->rb_right, new);
  150. }
  151. }
  152. static inline void rb_rotate_set_parents(RBNode *old, RBNode *new,
  153. RBRoot *root, RBColor color)
  154. {
  155. RBNode *parent = rb_parent(old);
  156. new->rb_parent_color = old->rb_parent_color;
  157. rb_set_parent_color(old, new, color);
  158. rb_change_child(old, new, parent, root);
  159. }
  160. static void rb_insert_augmented(RBNode *node, RBRoot *root,
  161. const RBAugmentCallbacks *augment)
  162. {
  163. RBNode *parent = rb_red_parent(node), *gparent, *tmp;
  164. while (true) {
  165. /*
  166. * Loop invariant: node is red.
  167. */
  168. if (unlikely(!parent)) {
  169. /*
  170. * The inserted node is root. Either this is the first node, or
  171. * we recursed at Case 1 below and are no longer violating 4).
  172. */
  173. rb_set_parent_color(node, NULL, RB_BLACK);
  174. break;
  175. }
  176. /*
  177. * If there is a black parent, we are done. Otherwise, take some
  178. * corrective action as, per 4), we don't want a red root or two
  179. * consecutive red nodes.
  180. */
  181. if (rb_is_black(parent)) {
  182. break;
  183. }
  184. gparent = rb_red_parent(parent);
  185. tmp = gparent->rb_right;
  186. if (parent != tmp) { /* parent == gparent->rb_left */
  187. if (tmp && rb_is_red(tmp)) {
  188. /*
  189. * Case 1 - node's uncle is red (color flips).
  190. *
  191. * G g
  192. * / \ / \
  193. * p u --> P U
  194. * / /
  195. * n n
  196. *
  197. * However, since g's parent might be red, and 4) does not
  198. * allow this, we need to recurse at g.
  199. */
  200. rb_set_parent_color(tmp, gparent, RB_BLACK);
  201. rb_set_parent_color(parent, gparent, RB_BLACK);
  202. node = gparent;
  203. parent = rb_parent(node);
  204. rb_set_parent_color(node, parent, RB_RED);
  205. continue;
  206. }
  207. tmp = parent->rb_right;
  208. if (node == tmp) {
  209. /*
  210. * Case 2 - node's uncle is black and node is
  211. * the parent's right child (left rotate at parent).
  212. *
  213. * G G
  214. * / \ / \
  215. * p U --> n U
  216. * \ /
  217. * n p
  218. *
  219. * This still leaves us in violation of 4), the
  220. * continuation into Case 3 will fix that.
  221. */
  222. tmp = node->rb_left;
  223. qatomic_set(&parent->rb_right, tmp);
  224. qatomic_set(&node->rb_left, parent);
  225. if (tmp) {
  226. rb_set_parent_color(tmp, parent, RB_BLACK);
  227. }
  228. rb_set_parent_color(parent, node, RB_RED);
  229. augment->rotate(parent, node);
  230. parent = node;
  231. tmp = node->rb_right;
  232. }
  233. /*
  234. * Case 3 - node's uncle is black and node is
  235. * the parent's left child (right rotate at gparent).
  236. *
  237. * G P
  238. * / \ / \
  239. * p U --> n g
  240. * / \
  241. * n U
  242. */
  243. qatomic_set(&gparent->rb_left, tmp); /* == parent->rb_right */
  244. qatomic_set(&parent->rb_right, gparent);
  245. if (tmp) {
  246. rb_set_parent_color(tmp, gparent, RB_BLACK);
  247. }
  248. rb_rotate_set_parents(gparent, parent, root, RB_RED);
  249. augment->rotate(gparent, parent);
  250. break;
  251. } else {
  252. tmp = gparent->rb_left;
  253. if (tmp && rb_is_red(tmp)) {
  254. /* Case 1 - color flips */
  255. rb_set_parent_color(tmp, gparent, RB_BLACK);
  256. rb_set_parent_color(parent, gparent, RB_BLACK);
  257. node = gparent;
  258. parent = rb_parent(node);
  259. rb_set_parent_color(node, parent, RB_RED);
  260. continue;
  261. }
  262. tmp = parent->rb_left;
  263. if (node == tmp) {
  264. /* Case 2 - right rotate at parent */
  265. tmp = node->rb_right;
  266. qatomic_set(&parent->rb_left, tmp);
  267. qatomic_set(&node->rb_right, parent);
  268. if (tmp) {
  269. rb_set_parent_color(tmp, parent, RB_BLACK);
  270. }
  271. rb_set_parent_color(parent, node, RB_RED);
  272. augment->rotate(parent, node);
  273. parent = node;
  274. tmp = node->rb_left;
  275. }
  276. /* Case 3 - left rotate at gparent */
  277. qatomic_set(&gparent->rb_right, tmp); /* == parent->rb_left */
  278. qatomic_set(&parent->rb_left, gparent);
  279. if (tmp) {
  280. rb_set_parent_color(tmp, gparent, RB_BLACK);
  281. }
  282. rb_rotate_set_parents(gparent, parent, root, RB_RED);
  283. augment->rotate(gparent, parent);
  284. break;
  285. }
  286. }
  287. }
  288. static void rb_insert_augmented_cached(RBNode *node,
  289. RBRootLeftCached *root, bool newleft,
  290. const RBAugmentCallbacks *augment)
  291. {
  292. if (newleft) {
  293. root->rb_leftmost = node;
  294. }
  295. rb_insert_augmented(node, &root->rb_root, augment);
  296. }
  297. static void rb_erase_color(RBNode *parent, RBRoot *root,
  298. const RBAugmentCallbacks *augment)
  299. {
  300. RBNode *node = NULL, *sibling, *tmp1, *tmp2;
  301. while (true) {
  302. /*
  303. * Loop invariants:
  304. * - node is black (or NULL on first iteration)
  305. * - node is not the root (parent is not NULL)
  306. * - All leaf paths going through parent and node have a
  307. * black node count that is 1 lower than other leaf paths.
  308. */
  309. sibling = parent->rb_right;
  310. if (node != sibling) { /* node == parent->rb_left */
  311. if (rb_is_red(sibling)) {
  312. /*
  313. * Case 1 - left rotate at parent
  314. *
  315. * P S
  316. * / \ / \
  317. * N s --> p Sr
  318. * / \ / \
  319. * Sl Sr N Sl
  320. */
  321. tmp1 = sibling->rb_left;
  322. qatomic_set(&parent->rb_right, tmp1);
  323. qatomic_set(&sibling->rb_left, parent);
  324. rb_set_parent_color(tmp1, parent, RB_BLACK);
  325. rb_rotate_set_parents(parent, sibling, root, RB_RED);
  326. augment->rotate(parent, sibling);
  327. sibling = tmp1;
  328. }
  329. tmp1 = sibling->rb_right;
  330. if (!tmp1 || rb_is_black(tmp1)) {
  331. tmp2 = sibling->rb_left;
  332. if (!tmp2 || rb_is_black(tmp2)) {
  333. /*
  334. * Case 2 - sibling color flip
  335. * (p could be either color here)
  336. *
  337. * (p) (p)
  338. * / \ / \
  339. * N S --> N s
  340. * / \ / \
  341. * Sl Sr Sl Sr
  342. *
  343. * This leaves us violating 5) which
  344. * can be fixed by flipping p to black
  345. * if it was red, or by recursing at p.
  346. * p is red when coming from Case 1.
  347. */
  348. rb_set_parent_color(sibling, parent, RB_RED);
  349. if (rb_is_red(parent)) {
  350. rb_set_black(parent);
  351. } else {
  352. node = parent;
  353. parent = rb_parent(node);
  354. if (parent) {
  355. continue;
  356. }
  357. }
  358. break;
  359. }
  360. /*
  361. * Case 3 - right rotate at sibling
  362. * (p could be either color here)
  363. *
  364. * (p) (p)
  365. * / \ / \
  366. * N S --> N sl
  367. * / \ \
  368. * sl Sr S
  369. * \
  370. * Sr
  371. *
  372. * Note: p might be red, and then bot
  373. * p and sl are red after rotation (which
  374. * breaks property 4). This is fixed in
  375. * Case 4 (in rb_rotate_set_parents()
  376. * which set sl the color of p
  377. * and set p RB_BLACK)
  378. *
  379. * (p) (sl)
  380. * / \ / \
  381. * N sl --> P S
  382. * \ / \
  383. * S N Sr
  384. * \
  385. * Sr
  386. */
  387. tmp1 = tmp2->rb_right;
  388. qatomic_set(&sibling->rb_left, tmp1);
  389. qatomic_set(&tmp2->rb_right, sibling);
  390. qatomic_set(&parent->rb_right, tmp2);
  391. if (tmp1) {
  392. rb_set_parent_color(tmp1, sibling, RB_BLACK);
  393. }
  394. augment->rotate(sibling, tmp2);
  395. tmp1 = sibling;
  396. sibling = tmp2;
  397. }
  398. /*
  399. * Case 4 - left rotate at parent + color flips
  400. * (p and sl could be either color here.
  401. * After rotation, p becomes black, s acquires
  402. * p's color, and sl keeps its color)
  403. *
  404. * (p) (s)
  405. * / \ / \
  406. * N S --> P Sr
  407. * / \ / \
  408. * (sl) sr N (sl)
  409. */
  410. tmp2 = sibling->rb_left;
  411. qatomic_set(&parent->rb_right, tmp2);
  412. qatomic_set(&sibling->rb_left, parent);
  413. rb_set_parent_color(tmp1, sibling, RB_BLACK);
  414. if (tmp2) {
  415. rb_set_parent(tmp2, parent);
  416. }
  417. rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
  418. augment->rotate(parent, sibling);
  419. break;
  420. } else {
  421. sibling = parent->rb_left;
  422. if (rb_is_red(sibling)) {
  423. /* Case 1 - right rotate at parent */
  424. tmp1 = sibling->rb_right;
  425. qatomic_set(&parent->rb_left, tmp1);
  426. qatomic_set(&sibling->rb_right, parent);
  427. rb_set_parent_color(tmp1, parent, RB_BLACK);
  428. rb_rotate_set_parents(parent, sibling, root, RB_RED);
  429. augment->rotate(parent, sibling);
  430. sibling = tmp1;
  431. }
  432. tmp1 = sibling->rb_left;
  433. if (!tmp1 || rb_is_black(tmp1)) {
  434. tmp2 = sibling->rb_right;
  435. if (!tmp2 || rb_is_black(tmp2)) {
  436. /* Case 2 - sibling color flip */
  437. rb_set_parent_color(sibling, parent, RB_RED);
  438. if (rb_is_red(parent)) {
  439. rb_set_black(parent);
  440. } else {
  441. node = parent;
  442. parent = rb_parent(node);
  443. if (parent) {
  444. continue;
  445. }
  446. }
  447. break;
  448. }
  449. /* Case 3 - left rotate at sibling */
  450. tmp1 = tmp2->rb_left;
  451. qatomic_set(&sibling->rb_right, tmp1);
  452. qatomic_set(&tmp2->rb_left, sibling);
  453. qatomic_set(&parent->rb_left, tmp2);
  454. if (tmp1) {
  455. rb_set_parent_color(tmp1, sibling, RB_BLACK);
  456. }
  457. augment->rotate(sibling, tmp2);
  458. tmp1 = sibling;
  459. sibling = tmp2;
  460. }
  461. /* Case 4 - right rotate at parent + color flips */
  462. tmp2 = sibling->rb_right;
  463. qatomic_set(&parent->rb_left, tmp2);
  464. qatomic_set(&sibling->rb_right, parent);
  465. rb_set_parent_color(tmp1, sibling, RB_BLACK);
  466. if (tmp2) {
  467. rb_set_parent(tmp2, parent);
  468. }
  469. rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
  470. augment->rotate(parent, sibling);
  471. break;
  472. }
  473. }
  474. }
  475. static void rb_erase_augmented(RBNode *node, RBRoot *root,
  476. const RBAugmentCallbacks *augment)
  477. {
  478. RBNode *child = node->rb_right;
  479. RBNode *tmp = node->rb_left;
  480. RBNode *parent, *rebalance;
  481. uintptr_t pc;
  482. if (!tmp) {
  483. /*
  484. * Case 1: node to erase has no more than 1 child (easy!)
  485. *
  486. * Note that if there is one child it must be red due to 5)
  487. * and node must be black due to 4). We adjust colors locally
  488. * so as to bypass rb_erase_color() later on.
  489. */
  490. pc = node->rb_parent_color;
  491. parent = rb_parent(node);
  492. rb_change_child(node, child, parent, root);
  493. if (child) {
  494. child->rb_parent_color = pc;
  495. rebalance = NULL;
  496. } else {
  497. rebalance = pc_is_black(pc) ? parent : NULL;
  498. }
  499. tmp = parent;
  500. } else if (!child) {
  501. /* Still case 1, but this time the child is node->rb_left */
  502. pc = node->rb_parent_color;
  503. parent = rb_parent(node);
  504. tmp->rb_parent_color = pc;
  505. rb_change_child(node, tmp, parent, root);
  506. rebalance = NULL;
  507. tmp = parent;
  508. } else {
  509. RBNode *successor = child, *child2;
  510. tmp = child->rb_left;
  511. if (!tmp) {
  512. /*
  513. * Case 2: node's successor is its right child
  514. *
  515. * (n) (s)
  516. * / \ / \
  517. * (x) (s) -> (x) (c)
  518. * \
  519. * (c)
  520. */
  521. parent = successor;
  522. child2 = successor->rb_right;
  523. augment->copy(node, successor);
  524. } else {
  525. /*
  526. * Case 3: node's successor is leftmost under
  527. * node's right child subtree
  528. *
  529. * (n) (s)
  530. * / \ / \
  531. * (x) (y) -> (x) (y)
  532. * / /
  533. * (p) (p)
  534. * / /
  535. * (s) (c)
  536. * \
  537. * (c)
  538. */
  539. do {
  540. parent = successor;
  541. successor = tmp;
  542. tmp = tmp->rb_left;
  543. } while (tmp);
  544. child2 = successor->rb_right;
  545. qatomic_set(&parent->rb_left, child2);
  546. qatomic_set(&successor->rb_right, child);
  547. rb_set_parent(child, successor);
  548. augment->copy(node, successor);
  549. augment->propagate(parent, successor);
  550. }
  551. tmp = node->rb_left;
  552. qatomic_set(&successor->rb_left, tmp);
  553. rb_set_parent(tmp, successor);
  554. pc = node->rb_parent_color;
  555. tmp = rb_parent(node);
  556. rb_change_child(node, successor, tmp, root);
  557. if (child2) {
  558. rb_set_parent_color(child2, parent, RB_BLACK);
  559. rebalance = NULL;
  560. } else {
  561. rebalance = rb_is_black(successor) ? parent : NULL;
  562. }
  563. successor->rb_parent_color = pc;
  564. tmp = successor;
  565. }
  566. augment->propagate(tmp, NULL);
  567. if (rebalance) {
  568. rb_erase_color(rebalance, root, augment);
  569. }
  570. }
  571. static void rb_erase_augmented_cached(RBNode *node, RBRootLeftCached *root,
  572. const RBAugmentCallbacks *augment)
  573. {
  574. if (root->rb_leftmost == node) {
  575. root->rb_leftmost = rb_next(node);
  576. }
  577. rb_erase_augmented(node, &root->rb_root, augment);
  578. }
  579. /*
  580. * Interval trees.
  581. *
  582. * Derived from lib/interval_tree.c and its dependencies,
  583. * especially include/linux/interval_tree_generic.h.
  584. */
  585. #define rb_to_itree(N) container_of(N, IntervalTreeNode, rb)
  586. static bool interval_tree_compute_max(IntervalTreeNode *node, bool exit)
  587. {
  588. IntervalTreeNode *child;
  589. uint64_t max = node->last;
  590. if (node->rb.rb_left) {
  591. child = rb_to_itree(node->rb.rb_left);
  592. if (child->subtree_last > max) {
  593. max = child->subtree_last;
  594. }
  595. }
  596. if (node->rb.rb_right) {
  597. child = rb_to_itree(node->rb.rb_right);
  598. if (child->subtree_last > max) {
  599. max = child->subtree_last;
  600. }
  601. }
  602. if (exit && node->subtree_last == max) {
  603. return true;
  604. }
  605. node->subtree_last = max;
  606. return false;
  607. }
  608. static void interval_tree_propagate(RBNode *rb, RBNode *stop)
  609. {
  610. while (rb != stop) {
  611. IntervalTreeNode *node = rb_to_itree(rb);
  612. if (interval_tree_compute_max(node, true)) {
  613. break;
  614. }
  615. rb = rb_parent(&node->rb);
  616. }
  617. }
  618. static void interval_tree_copy(RBNode *rb_old, RBNode *rb_new)
  619. {
  620. IntervalTreeNode *old = rb_to_itree(rb_old);
  621. IntervalTreeNode *new = rb_to_itree(rb_new);
  622. new->subtree_last = old->subtree_last;
  623. }
  624. static void interval_tree_rotate(RBNode *rb_old, RBNode *rb_new)
  625. {
  626. IntervalTreeNode *old = rb_to_itree(rb_old);
  627. IntervalTreeNode *new = rb_to_itree(rb_new);
  628. new->subtree_last = old->subtree_last;
  629. interval_tree_compute_max(old, false);
  630. }
  631. static const RBAugmentCallbacks interval_tree_augment = {
  632. .propagate = interval_tree_propagate,
  633. .copy = interval_tree_copy,
  634. .rotate = interval_tree_rotate,
  635. };
  636. /* Insert / remove interval nodes from the tree */
  637. void interval_tree_insert(IntervalTreeNode *node, IntervalTreeRoot *root)
  638. {
  639. RBNode **link = &root->rb_root.rb_node, *rb_parent = NULL;
  640. uint64_t start = node->start, last = node->last;
  641. IntervalTreeNode *parent;
  642. bool leftmost = true;
  643. while (*link) {
  644. rb_parent = *link;
  645. parent = rb_to_itree(rb_parent);
  646. if (parent->subtree_last < last) {
  647. parent->subtree_last = last;
  648. }
  649. if (start < parent->start) {
  650. link = &parent->rb.rb_left;
  651. } else {
  652. link = &parent->rb.rb_right;
  653. leftmost = false;
  654. }
  655. }
  656. node->subtree_last = last;
  657. rb_link_node(&node->rb, rb_parent, link);
  658. rb_insert_augmented_cached(&node->rb, root, leftmost,
  659. &interval_tree_augment);
  660. }
  661. void interval_tree_remove(IntervalTreeNode *node, IntervalTreeRoot *root)
  662. {
  663. rb_erase_augmented_cached(&node->rb, root, &interval_tree_augment);
  664. }
  665. /*
  666. * Iterate over intervals intersecting [start;last]
  667. *
  668. * Note that a node's interval intersects [start;last] iff:
  669. * Cond1: node->start <= last
  670. * and
  671. * Cond2: start <= node->last
  672. */
  673. static IntervalTreeNode *interval_tree_subtree_search(IntervalTreeNode *node,
  674. uint64_t start,
  675. uint64_t last)
  676. {
  677. while (true) {
  678. /*
  679. * Loop invariant: start <= node->subtree_last
  680. * (Cond2 is satisfied by one of the subtree nodes)
  681. */
  682. if (node->rb.rb_left) {
  683. IntervalTreeNode *left = rb_to_itree(node->rb.rb_left);
  684. if (start <= left->subtree_last) {
  685. /*
  686. * Some nodes in left subtree satisfy Cond2.
  687. * Iterate to find the leftmost such node N.
  688. * If it also satisfies Cond1, that's the
  689. * match we are looking for. Otherwise, there
  690. * is no matching interval as nodes to the
  691. * right of N can't satisfy Cond1 either.
  692. */
  693. node = left;
  694. continue;
  695. }
  696. }
  697. if (node->start <= last) { /* Cond1 */
  698. if (start <= node->last) { /* Cond2 */
  699. return node; /* node is leftmost match */
  700. }
  701. if (node->rb.rb_right) {
  702. node = rb_to_itree(node->rb.rb_right);
  703. if (start <= node->subtree_last) {
  704. continue;
  705. }
  706. }
  707. }
  708. return NULL; /* no match */
  709. }
  710. }
  711. IntervalTreeNode *interval_tree_iter_first(IntervalTreeRoot *root,
  712. uint64_t start, uint64_t last)
  713. {
  714. IntervalTreeNode *node, *leftmost;
  715. if (!root->rb_root.rb_node) {
  716. return NULL;
  717. }
  718. /*
  719. * Fastpath range intersection/overlap between A: [a0, a1] and
  720. * B: [b0, b1] is given by:
  721. *
  722. * a0 <= b1 && b0 <= a1
  723. *
  724. * ... where A holds the lock range and B holds the smallest
  725. * 'start' and largest 'last' in the tree. For the later, we
  726. * rely on the root node, which by augmented interval tree
  727. * property, holds the largest value in its last-in-subtree.
  728. * This allows mitigating some of the tree walk overhead for
  729. * for non-intersecting ranges, maintained and consulted in O(1).
  730. */
  731. node = rb_to_itree(root->rb_root.rb_node);
  732. if (node->subtree_last < start) {
  733. return NULL;
  734. }
  735. leftmost = rb_to_itree(root->rb_leftmost);
  736. if (leftmost->start > last) {
  737. return NULL;
  738. }
  739. return interval_tree_subtree_search(node, start, last);
  740. }
  741. IntervalTreeNode *interval_tree_iter_next(IntervalTreeNode *node,
  742. uint64_t start, uint64_t last)
  743. {
  744. RBNode *rb = node->rb.rb_right, *prev;
  745. while (true) {
  746. /*
  747. * Loop invariants:
  748. * Cond1: node->start <= last
  749. * rb == node->rb.rb_right
  750. *
  751. * First, search right subtree if suitable
  752. */
  753. if (rb) {
  754. IntervalTreeNode *right = rb_to_itree(rb);
  755. if (start <= right->subtree_last) {
  756. return interval_tree_subtree_search(right, start, last);
  757. }
  758. }
  759. /* Move up the tree until we come from a node's left child */
  760. do {
  761. rb = rb_parent(&node->rb);
  762. if (!rb) {
  763. return NULL;
  764. }
  765. prev = &node->rb;
  766. node = rb_to_itree(rb);
  767. rb = node->rb.rb_right;
  768. } while (prev == rb);
  769. /* Check if the node intersects [start;last] */
  770. if (last < node->start) { /* !Cond1 */
  771. return NULL;
  772. }
  773. if (start <= node->last) { /* Cond2 */
  774. return node;
  775. }
  776. }
  777. }
  778. /* Occasionally useful for calling from within the debugger. */
  779. #if 0
  780. static void debug_interval_tree_int(IntervalTreeNode *node,
  781. const char *dir, int level)
  782. {
  783. printf("%4d %*s %s [%" PRIu64 ",%" PRIu64 "] subtree_last:%" PRIu64 "\n",
  784. level, level + 1, dir, rb_is_red(&node->rb) ? "r" : "b",
  785. node->start, node->last, node->subtree_last);
  786. if (node->rb.rb_left) {
  787. debug_interval_tree_int(rb_to_itree(node->rb.rb_left), "<", level + 1);
  788. }
  789. if (node->rb.rb_right) {
  790. debug_interval_tree_int(rb_to_itree(node->rb.rb_right), ">", level + 1);
  791. }
  792. }
  793. void debug_interval_tree(IntervalTreeNode *node);
  794. void debug_interval_tree(IntervalTreeNode *node)
  795. {
  796. if (node) {
  797. debug_interval_tree_int(node, "*", 0);
  798. } else {
  799. printf("null\n");
  800. }
  801. }
  802. #endif