util/interval-tree.c

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