qemu/util/interval-tree.c
<<
>>
Prefs
   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
  53typedef enum RBColor
  54{
  55    RB_RED,
  56    RB_BLACK,
  57} RBColor;
  58
  59typedef 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
  65static inline uintptr_t rb_pc(const RBNode *n)
  66{
  67    return qatomic_read(&n->rb_parent_color);
  68}
  69
  70static inline void rb_set_pc(RBNode *n, uintptr_t pc)
  71{
  72    qatomic_set(&n->rb_parent_color, pc);
  73}
  74
  75static inline RBNode *pc_parent(uintptr_t pc)
  76{
  77    return (RBNode *)(pc & ~1);
  78}
  79
  80static inline RBNode *rb_parent(const RBNode *n)
  81{
  82    return pc_parent(rb_pc(n));
  83}
  84
  85static inline RBNode *rb_red_parent(const RBNode *n)
  86{
  87    return (RBNode *)rb_pc(n);
  88}
  89
  90static inline RBColor pc_color(uintptr_t pc)
  91{
  92    return (RBColor)(pc & 1);
  93}
  94
  95static inline bool pc_is_red(uintptr_t pc)
  96{
  97    return pc_color(pc) == RB_RED;
  98}
  99
 100static inline bool pc_is_black(uintptr_t pc)
 101{
 102    return !pc_is_red(pc);
 103}
 104
 105static inline RBColor rb_color(const RBNode *n)
 106{
 107    return pc_color(rb_pc(n));
 108}
 109
 110static inline bool rb_is_red(const RBNode *n)
 111{
 112    return pc_is_red(rb_pc(n));
 113}
 114
 115static inline bool rb_is_black(const RBNode *n)
 116{
 117    return pc_is_black(rb_pc(n));
 118}
 119
 120static inline void rb_set_black(RBNode *n)
 121{
 122    rb_set_pc(n, rb_pc(n) | RB_BLACK);
 123}
 124
 125static inline void rb_set_parent_color(RBNode *n, RBNode *p, RBColor color)
 126{
 127    rb_set_pc(n, (uintptr_t)p | color);
 128}
 129
 130static inline void rb_set_parent(RBNode *n, RBNode *p)
 131{
 132    rb_set_parent_color(n, p, rb_color(n));
 133}
 134
 135static 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
 147static 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
 178static 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
 190static 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
 201static 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
 338static 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
 348static 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
 528static 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
 633static 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
 652static 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
 676static 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
 687static 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
 695static 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
 704static 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 */
 711void 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
 739void 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
 753static 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
 795IntervalTreeNode *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
 830IntervalTreeNode *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
 875static 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
 890void debug_interval_tree(IntervalTreeNode *node);
 891void 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
 900