xref: /qemu/util/interval-tree.c (revision 47d1e982)
10d99d37aSRichard Henderson /* SPDX-License-Identifier: GPL-2.0-or-later */
20d99d37aSRichard Henderson 
30d99d37aSRichard Henderson #include "qemu/osdep.h"
40d99d37aSRichard Henderson #include "qemu/interval-tree.h"
50d99d37aSRichard Henderson #include "qemu/atomic.h"
60d99d37aSRichard Henderson 
70d99d37aSRichard Henderson /*
80d99d37aSRichard Henderson  * Red Black Trees.
90d99d37aSRichard Henderson  *
100d99d37aSRichard Henderson  * For now, don't expose Linux Red-Black Trees separately, but retain the
110d99d37aSRichard Henderson  * separate type definitions to keep the implementation sane, and allow
120d99d37aSRichard Henderson  * the possibility of separating them later.
130d99d37aSRichard Henderson  *
140d99d37aSRichard Henderson  * Derived from include/linux/rbtree_augmented.h and its dependencies.
150d99d37aSRichard Henderson  */
160d99d37aSRichard Henderson 
170d99d37aSRichard Henderson /*
180d99d37aSRichard Henderson  * red-black trees properties:  https://en.wikipedia.org/wiki/Rbtree
190d99d37aSRichard Henderson  *
200d99d37aSRichard Henderson  *  1) A node is either red or black
210d99d37aSRichard Henderson  *  2) The root is black
220d99d37aSRichard Henderson  *  3) All leaves (NULL) are black
230d99d37aSRichard Henderson  *  4) Both children of every red node are black
240d99d37aSRichard Henderson  *  5) Every simple path from root to leaves contains the same number
250d99d37aSRichard Henderson  *     of black nodes.
260d99d37aSRichard Henderson  *
270d99d37aSRichard Henderson  *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
280d99d37aSRichard Henderson  *  consecutive red nodes in a path and every red node is therefore followed by
290d99d37aSRichard Henderson  *  a black. So if B is the number of black nodes on every simple path (as per
300d99d37aSRichard Henderson  *  5), then the longest possible path due to 4 is 2B.
310d99d37aSRichard Henderson  *
320d99d37aSRichard Henderson  *  We shall indicate color with case, where black nodes are uppercase and red
330d99d37aSRichard Henderson  *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
340d99d37aSRichard Henderson  *  parentheses and have some accompanying text comment.
350d99d37aSRichard Henderson  *
360d99d37aSRichard Henderson  * Notes on lockless lookups:
370d99d37aSRichard Henderson  *
380d99d37aSRichard Henderson  * All stores to the tree structure (rb_left and rb_right) must be done using
390d99d37aSRichard Henderson  * WRITE_ONCE [qatomic_set for QEMU]. And we must not inadvertently cause
400d99d37aSRichard Henderson  * (temporary) loops in the tree structure as seen in program order.
410d99d37aSRichard Henderson  *
420d99d37aSRichard Henderson  * These two requirements will allow lockless iteration of the tree -- not
430d99d37aSRichard Henderson  * correct iteration mind you, tree rotations are not atomic so a lookup might
440d99d37aSRichard Henderson  * miss entire subtrees.
450d99d37aSRichard Henderson  *
460d99d37aSRichard Henderson  * But they do guarantee that any such traversal will only see valid elements
470d99d37aSRichard Henderson  * and that it will indeed complete -- does not get stuck in a loop.
480d99d37aSRichard Henderson  *
490d99d37aSRichard Henderson  * It also guarantees that if the lookup returns an element it is the 'correct'
500d99d37aSRichard Henderson  * one. But not returning an element does _NOT_ mean it's not present.
510d99d37aSRichard Henderson  */
520d99d37aSRichard Henderson 
530d99d37aSRichard Henderson typedef enum RBColor
540d99d37aSRichard Henderson {
550d99d37aSRichard Henderson     RB_RED,
560d99d37aSRichard Henderson     RB_BLACK,
570d99d37aSRichard Henderson } RBColor;
580d99d37aSRichard Henderson 
590d99d37aSRichard Henderson typedef struct RBAugmentCallbacks {
600d99d37aSRichard Henderson     void (*propagate)(RBNode *node, RBNode *stop);
610d99d37aSRichard Henderson     void (*copy)(RBNode *old, RBNode *new);
620d99d37aSRichard Henderson     void (*rotate)(RBNode *old, RBNode *new);
630d99d37aSRichard Henderson } RBAugmentCallbacks;
640d99d37aSRichard Henderson 
rb_pc(const RBNode * n)6579e29851SRichard Henderson static inline uintptr_t rb_pc(const RBNode *n)
6679e29851SRichard Henderson {
6779e29851SRichard Henderson     return qatomic_read(&n->rb_parent_color);
6879e29851SRichard Henderson }
6979e29851SRichard Henderson 
rb_set_pc(RBNode * n,uintptr_t pc)7079e29851SRichard Henderson static inline void rb_set_pc(RBNode *n, uintptr_t pc)
7179e29851SRichard Henderson {
7279e29851SRichard Henderson     qatomic_set(&n->rb_parent_color, pc);
7379e29851SRichard Henderson }
7479e29851SRichard Henderson 
pc_parent(uintptr_t pc)75d37a259fSRichard Henderson static inline RBNode *pc_parent(uintptr_t pc)
76d37a259fSRichard Henderson {
77d37a259fSRichard Henderson     return (RBNode *)(pc & ~1);
78d37a259fSRichard Henderson }
79d37a259fSRichard Henderson 
rb_parent(const RBNode * n)800d99d37aSRichard Henderson static inline RBNode *rb_parent(const RBNode *n)
810d99d37aSRichard Henderson {
8279e29851SRichard Henderson     return pc_parent(rb_pc(n));
830d99d37aSRichard Henderson }
840d99d37aSRichard Henderson 
rb_red_parent(const RBNode * n)850d99d37aSRichard Henderson static inline RBNode *rb_red_parent(const RBNode *n)
860d99d37aSRichard Henderson {
8779e29851SRichard Henderson     return (RBNode *)rb_pc(n);
880d99d37aSRichard Henderson }
890d99d37aSRichard Henderson 
pc_color(uintptr_t pc)900d99d37aSRichard Henderson static inline RBColor pc_color(uintptr_t pc)
910d99d37aSRichard Henderson {
920d99d37aSRichard Henderson     return (RBColor)(pc & 1);
930d99d37aSRichard Henderson }
940d99d37aSRichard Henderson 
pc_is_red(uintptr_t pc)950d99d37aSRichard Henderson static inline bool pc_is_red(uintptr_t pc)
960d99d37aSRichard Henderson {
970d99d37aSRichard Henderson     return pc_color(pc) == RB_RED;
980d99d37aSRichard Henderson }
990d99d37aSRichard Henderson 
pc_is_black(uintptr_t pc)1000d99d37aSRichard Henderson static inline bool pc_is_black(uintptr_t pc)
1010d99d37aSRichard Henderson {
1020d99d37aSRichard Henderson     return !pc_is_red(pc);
1030d99d37aSRichard Henderson }
1040d99d37aSRichard Henderson 
rb_color(const RBNode * n)1050d99d37aSRichard Henderson static inline RBColor rb_color(const RBNode *n)
1060d99d37aSRichard Henderson {
10779e29851SRichard Henderson     return pc_color(rb_pc(n));
1080d99d37aSRichard Henderson }
1090d99d37aSRichard Henderson 
rb_is_red(const RBNode * n)1100d99d37aSRichard Henderson static inline bool rb_is_red(const RBNode *n)
1110d99d37aSRichard Henderson {
11279e29851SRichard Henderson     return pc_is_red(rb_pc(n));
1130d99d37aSRichard Henderson }
1140d99d37aSRichard Henderson 
rb_is_black(const RBNode * n)1150d99d37aSRichard Henderson static inline bool rb_is_black(const RBNode *n)
1160d99d37aSRichard Henderson {
11779e29851SRichard Henderson     return pc_is_black(rb_pc(n));
1180d99d37aSRichard Henderson }
1190d99d37aSRichard Henderson 
rb_set_black(RBNode * n)1200d99d37aSRichard Henderson static inline void rb_set_black(RBNode *n)
1210d99d37aSRichard Henderson {
12279e29851SRichard Henderson     rb_set_pc(n, rb_pc(n) | RB_BLACK);
1230d99d37aSRichard Henderson }
1240d99d37aSRichard Henderson 
rb_set_parent_color(RBNode * n,RBNode * p,RBColor color)1250d99d37aSRichard Henderson static inline void rb_set_parent_color(RBNode *n, RBNode *p, RBColor color)
1260d99d37aSRichard Henderson {
12779e29851SRichard Henderson     rb_set_pc(n, (uintptr_t)p | color);
1280d99d37aSRichard Henderson }
1290d99d37aSRichard Henderson 
rb_set_parent(RBNode * n,RBNode * p)1300d99d37aSRichard Henderson static inline void rb_set_parent(RBNode *n, RBNode *p)
1310d99d37aSRichard Henderson {
1320d99d37aSRichard Henderson     rb_set_parent_color(n, p, rb_color(n));
1330d99d37aSRichard Henderson }
1340d99d37aSRichard Henderson 
rb_link_node(RBNode * node,RBNode * parent,RBNode ** rb_link)1350d99d37aSRichard Henderson static inline void rb_link_node(RBNode *node, RBNode *parent, RBNode **rb_link)
1360d99d37aSRichard Henderson {
1370d99d37aSRichard Henderson     node->rb_parent_color = (uintptr_t)parent;
1380d99d37aSRichard Henderson     node->rb_left = node->rb_right = NULL;
1390d99d37aSRichard Henderson 
1404c8baa02SRichard Henderson     /*
1414c8baa02SRichard Henderson      * Ensure that node is initialized before insertion,
1424c8baa02SRichard Henderson      * as viewed by a concurrent search.
1434c8baa02SRichard Henderson      */
1444c8baa02SRichard Henderson     qatomic_set_mb(rb_link, node);
1450d99d37aSRichard Henderson }
1460d99d37aSRichard Henderson 
rb_next(RBNode * node)1470d99d37aSRichard Henderson static RBNode *rb_next(RBNode *node)
1480d99d37aSRichard Henderson {
1490d99d37aSRichard Henderson     RBNode *parent;
1500d99d37aSRichard Henderson 
1510d99d37aSRichard Henderson     /* OMIT: if empty node, return null. */
1520d99d37aSRichard Henderson 
1530d99d37aSRichard Henderson     /*
1540d99d37aSRichard Henderson      * If we have a right-hand child, go down and then left as far as we can.
1550d99d37aSRichard Henderson      */
1560d99d37aSRichard Henderson     if (node->rb_right) {
1570d99d37aSRichard Henderson         node = node->rb_right;
1580d99d37aSRichard Henderson         while (node->rb_left) {
1590d99d37aSRichard Henderson             node = node->rb_left;
1600d99d37aSRichard Henderson         }
1610d99d37aSRichard Henderson         return node;
1620d99d37aSRichard Henderson     }
1630d99d37aSRichard Henderson 
1640d99d37aSRichard Henderson     /*
1650d99d37aSRichard Henderson      * No right-hand children. Everything down and left is smaller than us,
1660d99d37aSRichard Henderson      * so any 'next' node must be in the general direction of our parent.
1670d99d37aSRichard Henderson      * Go up the tree; any time the ancestor is a right-hand child of its
1680d99d37aSRichard Henderson      * parent, keep going up. First time it's a left-hand child of its
1690d99d37aSRichard Henderson      * parent, said parent is our 'next' node.
1700d99d37aSRichard Henderson      */
1710d99d37aSRichard Henderson     while ((parent = rb_parent(node)) && node == parent->rb_right) {
1720d99d37aSRichard Henderson         node = parent;
1730d99d37aSRichard Henderson     }
1740d99d37aSRichard Henderson 
1750d99d37aSRichard Henderson     return parent;
1760d99d37aSRichard Henderson }
1770d99d37aSRichard Henderson 
rb_change_child(RBNode * old,RBNode * new,RBNode * parent,RBRoot * root)1780d99d37aSRichard Henderson static inline void rb_change_child(RBNode *old, RBNode *new,
1790d99d37aSRichard Henderson                                    RBNode *parent, RBRoot *root)
1800d99d37aSRichard Henderson {
1810d99d37aSRichard Henderson     if (!parent) {
1820d99d37aSRichard Henderson         qatomic_set(&root->rb_node, new);
1830d99d37aSRichard Henderson     } else if (parent->rb_left == old) {
1840d99d37aSRichard Henderson         qatomic_set(&parent->rb_left, new);
1850d99d37aSRichard Henderson     } else {
1860d99d37aSRichard Henderson         qatomic_set(&parent->rb_right, new);
1870d99d37aSRichard Henderson     }
1880d99d37aSRichard Henderson }
1890d99d37aSRichard Henderson 
rb_rotate_set_parents(RBNode * old,RBNode * new,RBRoot * root,RBColor color)1900d99d37aSRichard Henderson static inline void rb_rotate_set_parents(RBNode *old, RBNode *new,
1910d99d37aSRichard Henderson                                          RBRoot *root, RBColor color)
1920d99d37aSRichard Henderson {
19379e29851SRichard Henderson     uintptr_t pc = rb_pc(old);
19479e29851SRichard Henderson     RBNode *parent = pc_parent(pc);
1950d99d37aSRichard Henderson 
19679e29851SRichard Henderson     rb_set_pc(new, pc);
1970d99d37aSRichard Henderson     rb_set_parent_color(old, new, color);
1980d99d37aSRichard Henderson     rb_change_child(old, new, parent, root);
1990d99d37aSRichard Henderson }
2000d99d37aSRichard Henderson 
rb_insert_augmented(RBNode * node,RBRoot * root,const RBAugmentCallbacks * augment)2010d99d37aSRichard Henderson static void rb_insert_augmented(RBNode *node, RBRoot *root,
2020d99d37aSRichard Henderson                                 const RBAugmentCallbacks *augment)
2030d99d37aSRichard Henderson {
2040d99d37aSRichard Henderson     RBNode *parent = rb_red_parent(node), *gparent, *tmp;
2050d99d37aSRichard Henderson 
2060d99d37aSRichard Henderson     while (true) {
2070d99d37aSRichard Henderson         /*
2080d99d37aSRichard Henderson          * Loop invariant: node is red.
2090d99d37aSRichard Henderson          */
2100d99d37aSRichard Henderson         if (unlikely(!parent)) {
2110d99d37aSRichard Henderson             /*
2120d99d37aSRichard Henderson              * The inserted node is root. Either this is the first node, or
2130d99d37aSRichard Henderson              * we recursed at Case 1 below and are no longer violating 4).
2140d99d37aSRichard Henderson              */
2150d99d37aSRichard Henderson             rb_set_parent_color(node, NULL, RB_BLACK);
2160d99d37aSRichard Henderson             break;
2170d99d37aSRichard Henderson         }
2180d99d37aSRichard Henderson 
2190d99d37aSRichard Henderson         /*
2200d99d37aSRichard Henderson          * If there is a black parent, we are done.  Otherwise, take some
2210d99d37aSRichard Henderson          * corrective action as, per 4), we don't want a red root or two
2220d99d37aSRichard Henderson          * consecutive red nodes.
2230d99d37aSRichard Henderson          */
2240d99d37aSRichard Henderson         if (rb_is_black(parent)) {
2250d99d37aSRichard Henderson             break;
2260d99d37aSRichard Henderson         }
2270d99d37aSRichard Henderson 
2280d99d37aSRichard Henderson         gparent = rb_red_parent(parent);
2290d99d37aSRichard Henderson 
2300d99d37aSRichard Henderson         tmp = gparent->rb_right;
2310d99d37aSRichard Henderson         if (parent != tmp) {    /* parent == gparent->rb_left */
2320d99d37aSRichard Henderson             if (tmp && rb_is_red(tmp)) {
2330d99d37aSRichard Henderson                 /*
2340d99d37aSRichard Henderson                  * Case 1 - node's uncle is red (color flips).
2350d99d37aSRichard Henderson                  *
2360d99d37aSRichard Henderson                  *       G            g
2370d99d37aSRichard Henderson                  *      / \          / \
2380d99d37aSRichard Henderson                  *     p   u  -->   P   U
2390d99d37aSRichard Henderson                  *    /            /
2400d99d37aSRichard Henderson                  *   n            n
2410d99d37aSRichard Henderson                  *
2420d99d37aSRichard Henderson                  * However, since g's parent might be red, and 4) does not
2430d99d37aSRichard Henderson                  * allow this, we need to recurse at g.
2440d99d37aSRichard Henderson                  */
2450d99d37aSRichard Henderson                 rb_set_parent_color(tmp, gparent, RB_BLACK);
2460d99d37aSRichard Henderson                 rb_set_parent_color(parent, gparent, RB_BLACK);
2470d99d37aSRichard Henderson                 node = gparent;
2480d99d37aSRichard Henderson                 parent = rb_parent(node);
2490d99d37aSRichard Henderson                 rb_set_parent_color(node, parent, RB_RED);
2500d99d37aSRichard Henderson                 continue;
2510d99d37aSRichard Henderson             }
2520d99d37aSRichard Henderson 
2530d99d37aSRichard Henderson             tmp = parent->rb_right;
2540d99d37aSRichard Henderson             if (node == tmp) {
2550d99d37aSRichard Henderson                 /*
2560d99d37aSRichard Henderson                  * Case 2 - node's uncle is black and node is
2570d99d37aSRichard Henderson                  * the parent's right child (left rotate at parent).
2580d99d37aSRichard Henderson                  *
2590d99d37aSRichard Henderson                  *      G             G
2600d99d37aSRichard Henderson                  *     / \           / \
2610d99d37aSRichard Henderson                  *    p   U  -->    n   U
2620d99d37aSRichard Henderson                  *     \           /
2630d99d37aSRichard Henderson                  *      n         p
2640d99d37aSRichard Henderson                  *
2650d99d37aSRichard Henderson                  * This still leaves us in violation of 4), the
2660d99d37aSRichard Henderson                  * continuation into Case 3 will fix that.
2670d99d37aSRichard Henderson                  */
2680d99d37aSRichard Henderson                 tmp = node->rb_left;
2690d99d37aSRichard Henderson                 qatomic_set(&parent->rb_right, tmp);
2700d99d37aSRichard Henderson                 qatomic_set(&node->rb_left, parent);
2710d99d37aSRichard Henderson                 if (tmp) {
2720d99d37aSRichard Henderson                     rb_set_parent_color(tmp, parent, RB_BLACK);
2730d99d37aSRichard Henderson                 }
2740d99d37aSRichard Henderson                 rb_set_parent_color(parent, node, RB_RED);
2750d99d37aSRichard Henderson                 augment->rotate(parent, node);
2760d99d37aSRichard Henderson                 parent = node;
2770d99d37aSRichard Henderson                 tmp = node->rb_right;
2780d99d37aSRichard Henderson             }
2790d99d37aSRichard Henderson 
2800d99d37aSRichard Henderson             /*
2810d99d37aSRichard Henderson              * Case 3 - node's uncle is black and node is
2820d99d37aSRichard Henderson              * the parent's left child (right rotate at gparent).
2830d99d37aSRichard Henderson              *
2840d99d37aSRichard Henderson              *        G           P
2850d99d37aSRichard Henderson              *       / \         / \
2860d99d37aSRichard Henderson              *      p   U  -->  n   g
2870d99d37aSRichard Henderson              *     /                 \
2880d99d37aSRichard Henderson              *    n                   U
2890d99d37aSRichard Henderson              */
2900d99d37aSRichard Henderson             qatomic_set(&gparent->rb_left, tmp); /* == parent->rb_right */
2910d99d37aSRichard Henderson             qatomic_set(&parent->rb_right, gparent);
2920d99d37aSRichard Henderson             if (tmp) {
2930d99d37aSRichard Henderson                 rb_set_parent_color(tmp, gparent, RB_BLACK);
2940d99d37aSRichard Henderson             }
2950d99d37aSRichard Henderson             rb_rotate_set_parents(gparent, parent, root, RB_RED);
2960d99d37aSRichard Henderson             augment->rotate(gparent, parent);
2970d99d37aSRichard Henderson             break;
2980d99d37aSRichard Henderson         } else {
2990d99d37aSRichard Henderson             tmp = gparent->rb_left;
3000d99d37aSRichard Henderson             if (tmp && rb_is_red(tmp)) {
3010d99d37aSRichard Henderson                 /* Case 1 - color flips */
3020d99d37aSRichard Henderson                 rb_set_parent_color(tmp, gparent, RB_BLACK);
3030d99d37aSRichard Henderson                 rb_set_parent_color(parent, gparent, RB_BLACK);
3040d99d37aSRichard Henderson                 node = gparent;
3050d99d37aSRichard Henderson                 parent = rb_parent(node);
3060d99d37aSRichard Henderson                 rb_set_parent_color(node, parent, RB_RED);
3070d99d37aSRichard Henderson                 continue;
3080d99d37aSRichard Henderson             }
3090d99d37aSRichard Henderson 
3100d99d37aSRichard Henderson             tmp = parent->rb_left;
3110d99d37aSRichard Henderson             if (node == tmp) {
3120d99d37aSRichard Henderson                 /* Case 2 - right rotate at parent */
3130d99d37aSRichard Henderson                 tmp = node->rb_right;
3140d99d37aSRichard Henderson                 qatomic_set(&parent->rb_left, tmp);
3150d99d37aSRichard Henderson                 qatomic_set(&node->rb_right, parent);
3160d99d37aSRichard Henderson                 if (tmp) {
3170d99d37aSRichard Henderson                     rb_set_parent_color(tmp, parent, RB_BLACK);
3180d99d37aSRichard Henderson                 }
3190d99d37aSRichard Henderson                 rb_set_parent_color(parent, node, RB_RED);
3200d99d37aSRichard Henderson                 augment->rotate(parent, node);
3210d99d37aSRichard Henderson                 parent = node;
3220d99d37aSRichard Henderson                 tmp = node->rb_left;
3230d99d37aSRichard Henderson             }
3240d99d37aSRichard Henderson 
3250d99d37aSRichard Henderson             /* Case 3 - left rotate at gparent */
3260d99d37aSRichard Henderson             qatomic_set(&gparent->rb_right, tmp); /* == parent->rb_left */
3270d99d37aSRichard Henderson             qatomic_set(&parent->rb_left, gparent);
3280d99d37aSRichard Henderson             if (tmp) {
3290d99d37aSRichard Henderson                 rb_set_parent_color(tmp, gparent, RB_BLACK);
3300d99d37aSRichard Henderson             }
3310d99d37aSRichard Henderson             rb_rotate_set_parents(gparent, parent, root, RB_RED);
3320d99d37aSRichard Henderson             augment->rotate(gparent, parent);
3330d99d37aSRichard Henderson             break;
3340d99d37aSRichard Henderson         }
3350d99d37aSRichard Henderson     }
3360d99d37aSRichard Henderson }
3370d99d37aSRichard Henderson 
rb_insert_augmented_cached(RBNode * node,RBRootLeftCached * root,bool newleft,const RBAugmentCallbacks * augment)3380d99d37aSRichard Henderson static void rb_insert_augmented_cached(RBNode *node,
3390d99d37aSRichard Henderson                                        RBRootLeftCached *root, bool newleft,
3400d99d37aSRichard Henderson                                        const RBAugmentCallbacks *augment)
3410d99d37aSRichard Henderson {
3420d99d37aSRichard Henderson     if (newleft) {
3430d99d37aSRichard Henderson         root->rb_leftmost = node;
3440d99d37aSRichard Henderson     }
3450d99d37aSRichard Henderson     rb_insert_augmented(node, &root->rb_root, augment);
3460d99d37aSRichard Henderson }
3470d99d37aSRichard Henderson 
rb_erase_color(RBNode * parent,RBRoot * root,const RBAugmentCallbacks * augment)3480d99d37aSRichard Henderson static void rb_erase_color(RBNode *parent, RBRoot *root,
3490d99d37aSRichard Henderson                            const RBAugmentCallbacks *augment)
3500d99d37aSRichard Henderson {
3510d99d37aSRichard Henderson     RBNode *node = NULL, *sibling, *tmp1, *tmp2;
3520d99d37aSRichard Henderson 
3530d99d37aSRichard Henderson     while (true) {
3540d99d37aSRichard Henderson         /*
3550d99d37aSRichard Henderson          * Loop invariants:
3560d99d37aSRichard Henderson          * - node is black (or NULL on first iteration)
3570d99d37aSRichard Henderson          * - node is not the root (parent is not NULL)
3580d99d37aSRichard Henderson          * - All leaf paths going through parent and node have a
3590d99d37aSRichard Henderson          *   black node count that is 1 lower than other leaf paths.
3600d99d37aSRichard Henderson          */
3610d99d37aSRichard Henderson         sibling = parent->rb_right;
3620d99d37aSRichard Henderson         if (node != sibling) {  /* node == parent->rb_left */
3630d99d37aSRichard Henderson             if (rb_is_red(sibling)) {
3640d99d37aSRichard Henderson                 /*
3650d99d37aSRichard Henderson                  * Case 1 - left rotate at parent
3660d99d37aSRichard Henderson                  *
3670d99d37aSRichard Henderson                  *     P               S
3680d99d37aSRichard Henderson                  *    / \             / \
3690d99d37aSRichard Henderson                  *   N   s    -->    p   Sr
3700d99d37aSRichard Henderson                  *      / \         / \
3710d99d37aSRichard Henderson                  *     Sl  Sr      N   Sl
3720d99d37aSRichard Henderson                  */
3730d99d37aSRichard Henderson                 tmp1 = sibling->rb_left;
3740d99d37aSRichard Henderson                 qatomic_set(&parent->rb_right, tmp1);
3750d99d37aSRichard Henderson                 qatomic_set(&sibling->rb_left, parent);
3760d99d37aSRichard Henderson                 rb_set_parent_color(tmp1, parent, RB_BLACK);
3770d99d37aSRichard Henderson                 rb_rotate_set_parents(parent, sibling, root, RB_RED);
3780d99d37aSRichard Henderson                 augment->rotate(parent, sibling);
3790d99d37aSRichard Henderson                 sibling = tmp1;
3800d99d37aSRichard Henderson             }
3810d99d37aSRichard Henderson             tmp1 = sibling->rb_right;
3820d99d37aSRichard Henderson             if (!tmp1 || rb_is_black(tmp1)) {
3830d99d37aSRichard Henderson                 tmp2 = sibling->rb_left;
3840d99d37aSRichard Henderson                 if (!tmp2 || rb_is_black(tmp2)) {
3850d99d37aSRichard Henderson                     /*
3860d99d37aSRichard Henderson                      * Case 2 - sibling color flip
3870d99d37aSRichard Henderson                      * (p could be either color here)
3880d99d37aSRichard Henderson                      *
3890d99d37aSRichard Henderson                      *    (p)           (p)
3900d99d37aSRichard Henderson                      *    / \           / \
3910d99d37aSRichard Henderson                      *   N   S    -->  N   s
3920d99d37aSRichard Henderson                      *      / \           / \
3930d99d37aSRichard Henderson                      *     Sl  Sr        Sl  Sr
3940d99d37aSRichard Henderson                      *
3950d99d37aSRichard Henderson                      * This leaves us violating 5) which
3960d99d37aSRichard Henderson                      * can be fixed by flipping p to black
3970d99d37aSRichard Henderson                      * if it was red, or by recursing at p.
3980d99d37aSRichard Henderson                      * p is red when coming from Case 1.
3990d99d37aSRichard Henderson                      */
4000d99d37aSRichard Henderson                     rb_set_parent_color(sibling, parent, RB_RED);
4010d99d37aSRichard Henderson                     if (rb_is_red(parent)) {
4020d99d37aSRichard Henderson                         rb_set_black(parent);
4030d99d37aSRichard Henderson                     } else {
4040d99d37aSRichard Henderson                         node = parent;
4050d99d37aSRichard Henderson                         parent = rb_parent(node);
4060d99d37aSRichard Henderson                         if (parent) {
4070d99d37aSRichard Henderson                             continue;
4080d99d37aSRichard Henderson                         }
4090d99d37aSRichard Henderson                     }
4100d99d37aSRichard Henderson                     break;
4110d99d37aSRichard Henderson                 }
4120d99d37aSRichard Henderson                 /*
4130d99d37aSRichard Henderson                  * Case 3 - right rotate at sibling
4140d99d37aSRichard Henderson                  * (p could be either color here)
4150d99d37aSRichard Henderson                  *
4160d99d37aSRichard Henderson                  *   (p)           (p)
4170d99d37aSRichard Henderson                  *   / \           / \
4180d99d37aSRichard Henderson                  *  N   S    -->  N   sl
4190d99d37aSRichard Henderson                  *     / \             \
4200d99d37aSRichard Henderson                  *    sl  Sr            S
4210d99d37aSRichard Henderson                  *                       \
4220d99d37aSRichard Henderson                  *                        Sr
4230d99d37aSRichard Henderson                  *
4240d99d37aSRichard Henderson                  * Note: p might be red, and then bot
4250d99d37aSRichard Henderson                  * p and sl are red after rotation (which
4260d99d37aSRichard Henderson                  * breaks property 4). This is fixed in
4270d99d37aSRichard Henderson                  * Case 4 (in rb_rotate_set_parents()
4280d99d37aSRichard Henderson                  *         which set sl the color of p
4290d99d37aSRichard Henderson                  *         and set p RB_BLACK)
4300d99d37aSRichard Henderson                  *
4310d99d37aSRichard Henderson                  *   (p)            (sl)
4320d99d37aSRichard Henderson                  *   / \            /  \
4330d99d37aSRichard Henderson                  *  N   sl   -->   P    S
4340d99d37aSRichard Henderson                  *       \        /      \
4350d99d37aSRichard Henderson                  *        S      N        Sr
4360d99d37aSRichard Henderson                  *         \
4370d99d37aSRichard Henderson                  *          Sr
4380d99d37aSRichard Henderson                  */
4390d99d37aSRichard Henderson                 tmp1 = tmp2->rb_right;
4400d99d37aSRichard Henderson                 qatomic_set(&sibling->rb_left, tmp1);
4410d99d37aSRichard Henderson                 qatomic_set(&tmp2->rb_right, sibling);
4420d99d37aSRichard Henderson                 qatomic_set(&parent->rb_right, tmp2);
4430d99d37aSRichard Henderson                 if (tmp1) {
4440d99d37aSRichard Henderson                     rb_set_parent_color(tmp1, sibling, RB_BLACK);
4450d99d37aSRichard Henderson                 }
4460d99d37aSRichard Henderson                 augment->rotate(sibling, tmp2);
4470d99d37aSRichard Henderson                 tmp1 = sibling;
4480d99d37aSRichard Henderson                 sibling = tmp2;
4490d99d37aSRichard Henderson             }
4500d99d37aSRichard Henderson             /*
4510d99d37aSRichard Henderson              * Case 4 - left rotate at parent + color flips
4520d99d37aSRichard Henderson              * (p and sl could be either color here.
4530d99d37aSRichard Henderson              *  After rotation, p becomes black, s acquires
4540d99d37aSRichard Henderson              *  p's color, and sl keeps its color)
4550d99d37aSRichard Henderson              *
4560d99d37aSRichard Henderson              *      (p)             (s)
4570d99d37aSRichard Henderson              *      / \             / \
4580d99d37aSRichard Henderson              *     N   S     -->   P   Sr
4590d99d37aSRichard Henderson              *        / \         / \
4600d99d37aSRichard Henderson              *      (sl) sr      N  (sl)
4610d99d37aSRichard Henderson              */
4620d99d37aSRichard Henderson             tmp2 = sibling->rb_left;
4630d99d37aSRichard Henderson             qatomic_set(&parent->rb_right, tmp2);
4640d99d37aSRichard Henderson             qatomic_set(&sibling->rb_left, parent);
4650d99d37aSRichard Henderson             rb_set_parent_color(tmp1, sibling, RB_BLACK);
4660d99d37aSRichard Henderson             if (tmp2) {
4670d99d37aSRichard Henderson                 rb_set_parent(tmp2, parent);
4680d99d37aSRichard Henderson             }
4690d99d37aSRichard Henderson             rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
4700d99d37aSRichard Henderson             augment->rotate(parent, sibling);
4710d99d37aSRichard Henderson             break;
4720d99d37aSRichard Henderson         } else {
4730d99d37aSRichard Henderson             sibling = parent->rb_left;
4740d99d37aSRichard Henderson             if (rb_is_red(sibling)) {
4750d99d37aSRichard Henderson                 /* Case 1 - right rotate at parent */
4760d99d37aSRichard Henderson                 tmp1 = sibling->rb_right;
4770d99d37aSRichard Henderson                 qatomic_set(&parent->rb_left, tmp1);
4780d99d37aSRichard Henderson                 qatomic_set(&sibling->rb_right, parent);
4790d99d37aSRichard Henderson                 rb_set_parent_color(tmp1, parent, RB_BLACK);
4800d99d37aSRichard Henderson                 rb_rotate_set_parents(parent, sibling, root, RB_RED);
4810d99d37aSRichard Henderson                 augment->rotate(parent, sibling);
4820d99d37aSRichard Henderson                 sibling = tmp1;
4830d99d37aSRichard Henderson             }
4840d99d37aSRichard Henderson             tmp1 = sibling->rb_left;
4850d99d37aSRichard Henderson             if (!tmp1 || rb_is_black(tmp1)) {
4860d99d37aSRichard Henderson                 tmp2 = sibling->rb_right;
4870d99d37aSRichard Henderson                 if (!tmp2 || rb_is_black(tmp2)) {
4880d99d37aSRichard Henderson                     /* Case 2 - sibling color flip */
4890d99d37aSRichard Henderson                     rb_set_parent_color(sibling, parent, RB_RED);
4900d99d37aSRichard Henderson                     if (rb_is_red(parent)) {
4910d99d37aSRichard Henderson                         rb_set_black(parent);
4920d99d37aSRichard Henderson                     } else {
4930d99d37aSRichard Henderson                         node = parent;
4940d99d37aSRichard Henderson                         parent = rb_parent(node);
4950d99d37aSRichard Henderson                         if (parent) {
4960d99d37aSRichard Henderson                             continue;
4970d99d37aSRichard Henderson                         }
4980d99d37aSRichard Henderson                     }
4990d99d37aSRichard Henderson                     break;
5000d99d37aSRichard Henderson                 }
5010d99d37aSRichard Henderson                 /* Case 3 - left rotate at sibling */
5020d99d37aSRichard Henderson                 tmp1 = tmp2->rb_left;
5030d99d37aSRichard Henderson                 qatomic_set(&sibling->rb_right, tmp1);
5040d99d37aSRichard Henderson                 qatomic_set(&tmp2->rb_left, sibling);
5050d99d37aSRichard Henderson                 qatomic_set(&parent->rb_left, tmp2);
5060d99d37aSRichard Henderson                 if (tmp1) {
5070d99d37aSRichard Henderson                     rb_set_parent_color(tmp1, sibling, RB_BLACK);
5080d99d37aSRichard Henderson                 }
5090d99d37aSRichard Henderson                 augment->rotate(sibling, tmp2);
5100d99d37aSRichard Henderson                 tmp1 = sibling;
5110d99d37aSRichard Henderson                 sibling = tmp2;
5120d99d37aSRichard Henderson             }
5130d99d37aSRichard Henderson             /* Case 4 - right rotate at parent + color flips */
5140d99d37aSRichard Henderson             tmp2 = sibling->rb_right;
5150d99d37aSRichard Henderson             qatomic_set(&parent->rb_left, tmp2);
5160d99d37aSRichard Henderson             qatomic_set(&sibling->rb_right, parent);
5170d99d37aSRichard Henderson             rb_set_parent_color(tmp1, sibling, RB_BLACK);
5180d99d37aSRichard Henderson             if (tmp2) {
5190d99d37aSRichard Henderson                 rb_set_parent(tmp2, parent);
5200d99d37aSRichard Henderson             }
5210d99d37aSRichard Henderson             rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
5220d99d37aSRichard Henderson             augment->rotate(parent, sibling);
5230d99d37aSRichard Henderson             break;
5240d99d37aSRichard Henderson         }
5250d99d37aSRichard Henderson     }
5260d99d37aSRichard Henderson }
5270d99d37aSRichard Henderson 
rb_erase_augmented(RBNode * node,RBRoot * root,const RBAugmentCallbacks * augment)5280d99d37aSRichard Henderson static void rb_erase_augmented(RBNode *node, RBRoot *root,
5290d99d37aSRichard Henderson                                const RBAugmentCallbacks *augment)
5300d99d37aSRichard Henderson {
5310d99d37aSRichard Henderson     RBNode *child = node->rb_right;
5320d99d37aSRichard Henderson     RBNode *tmp = node->rb_left;
5330d99d37aSRichard Henderson     RBNode *parent, *rebalance;
5340d99d37aSRichard Henderson     uintptr_t pc;
5350d99d37aSRichard Henderson 
5360d99d37aSRichard Henderson     if (!tmp) {
5370d99d37aSRichard Henderson         /*
5380d99d37aSRichard Henderson          * Case 1: node to erase has no more than 1 child (easy!)
5390d99d37aSRichard Henderson          *
5400d99d37aSRichard Henderson          * Note that if there is one child it must be red due to 5)
5410d99d37aSRichard Henderson          * and node must be black due to 4). We adjust colors locally
5420d99d37aSRichard Henderson          * so as to bypass rb_erase_color() later on.
5430d99d37aSRichard Henderson          */
54479e29851SRichard Henderson         pc = rb_pc(node);
545d37a259fSRichard Henderson         parent = pc_parent(pc);
5460d99d37aSRichard Henderson         rb_change_child(node, child, parent, root);
5470d99d37aSRichard Henderson         if (child) {
54879e29851SRichard Henderson             rb_set_pc(child, pc);
5490d99d37aSRichard Henderson             rebalance = NULL;
5500d99d37aSRichard Henderson         } else {
5510d99d37aSRichard Henderson             rebalance = pc_is_black(pc) ? parent : NULL;
5520d99d37aSRichard Henderson         }
5530d99d37aSRichard Henderson         tmp = parent;
5540d99d37aSRichard Henderson     } else if (!child) {
5550d99d37aSRichard Henderson         /* Still case 1, but this time the child is node->rb_left */
55679e29851SRichard Henderson         pc = rb_pc(node);
557d37a259fSRichard Henderson         parent = pc_parent(pc);
55879e29851SRichard Henderson         rb_set_pc(tmp, pc);
5590d99d37aSRichard Henderson         rb_change_child(node, tmp, parent, root);
5600d99d37aSRichard Henderson         rebalance = NULL;
5610d99d37aSRichard Henderson         tmp = parent;
5620d99d37aSRichard Henderson     } else {
5630d99d37aSRichard Henderson         RBNode *successor = child, *child2;
5640d99d37aSRichard Henderson         tmp = child->rb_left;
5650d99d37aSRichard Henderson         if (!tmp) {
5660d99d37aSRichard Henderson             /*
5670d99d37aSRichard Henderson              * Case 2: node's successor is its right child
5680d99d37aSRichard Henderson              *
5690d99d37aSRichard Henderson              *    (n)          (s)
5700d99d37aSRichard Henderson              *    / \          / \
5710d99d37aSRichard Henderson              *  (x) (s)  ->  (x) (c)
5720d99d37aSRichard Henderson              *        \
5730d99d37aSRichard Henderson              *        (c)
5740d99d37aSRichard Henderson              */
5750d99d37aSRichard Henderson             parent = successor;
5760d99d37aSRichard Henderson             child2 = successor->rb_right;
5770d99d37aSRichard Henderson 
5780d99d37aSRichard Henderson             augment->copy(node, successor);
5790d99d37aSRichard Henderson         } else {
5800d99d37aSRichard Henderson             /*
5810d99d37aSRichard Henderson              * Case 3: node's successor is leftmost under
5820d99d37aSRichard Henderson              * node's right child subtree
5830d99d37aSRichard Henderson              *
5840d99d37aSRichard Henderson              *    (n)          (s)
5850d99d37aSRichard Henderson              *    / \          / \
5860d99d37aSRichard Henderson              *  (x) (y)  ->  (x) (y)
5870d99d37aSRichard Henderson              *      /            /
5880d99d37aSRichard Henderson              *    (p)          (p)
5890d99d37aSRichard Henderson              *    /            /
5900d99d37aSRichard Henderson              *  (s)          (c)
5910d99d37aSRichard Henderson              *    \
5920d99d37aSRichard Henderson              *    (c)
5930d99d37aSRichard Henderson              */
5940d99d37aSRichard Henderson             do {
5950d99d37aSRichard Henderson                 parent = successor;
5960d99d37aSRichard Henderson                 successor = tmp;
5970d99d37aSRichard Henderson                 tmp = tmp->rb_left;
5980d99d37aSRichard Henderson             } while (tmp);
5990d99d37aSRichard Henderson             child2 = successor->rb_right;
6000d99d37aSRichard Henderson             qatomic_set(&parent->rb_left, child2);
6010d99d37aSRichard Henderson             qatomic_set(&successor->rb_right, child);
6020d99d37aSRichard Henderson             rb_set_parent(child, successor);
6030d99d37aSRichard Henderson 
6040d99d37aSRichard Henderson             augment->copy(node, successor);
6050d99d37aSRichard Henderson             augment->propagate(parent, successor);
6060d99d37aSRichard Henderson         }
6070d99d37aSRichard Henderson 
6080d99d37aSRichard Henderson         tmp = node->rb_left;
6090d99d37aSRichard Henderson         qatomic_set(&successor->rb_left, tmp);
6100d99d37aSRichard Henderson         rb_set_parent(tmp, successor);
6110d99d37aSRichard Henderson 
61279e29851SRichard Henderson         pc = rb_pc(node);
613d37a259fSRichard Henderson         tmp = pc_parent(pc);
6140d99d37aSRichard Henderson         rb_change_child(node, successor, tmp, root);
6150d99d37aSRichard Henderson 
6160d99d37aSRichard Henderson         if (child2) {
6170d99d37aSRichard Henderson             rb_set_parent_color(child2, parent, RB_BLACK);
6180d99d37aSRichard Henderson             rebalance = NULL;
6190d99d37aSRichard Henderson         } else {
6200d99d37aSRichard Henderson             rebalance = rb_is_black(successor) ? parent : NULL;
6210d99d37aSRichard Henderson         }
62279e29851SRichard Henderson         rb_set_pc(successor, pc);
6230d99d37aSRichard Henderson         tmp = successor;
6240d99d37aSRichard Henderson     }
6250d99d37aSRichard Henderson 
6260d99d37aSRichard Henderson     augment->propagate(tmp, NULL);
6270d99d37aSRichard Henderson 
6280d99d37aSRichard Henderson     if (rebalance) {
6290d99d37aSRichard Henderson         rb_erase_color(rebalance, root, augment);
6300d99d37aSRichard Henderson     }
6310d99d37aSRichard Henderson }
6320d99d37aSRichard Henderson 
rb_erase_augmented_cached(RBNode * node,RBRootLeftCached * root,const RBAugmentCallbacks * augment)6330d99d37aSRichard Henderson static void rb_erase_augmented_cached(RBNode *node, RBRootLeftCached *root,
6340d99d37aSRichard Henderson                                       const RBAugmentCallbacks *augment)
6350d99d37aSRichard Henderson {
6360d99d37aSRichard Henderson     if (root->rb_leftmost == node) {
6370d99d37aSRichard Henderson         root->rb_leftmost = rb_next(node);
6380d99d37aSRichard Henderson     }
6390d99d37aSRichard Henderson     rb_erase_augmented(node, &root->rb_root, augment);
6400d99d37aSRichard Henderson }
6410d99d37aSRichard Henderson 
6420d99d37aSRichard Henderson 
6430d99d37aSRichard Henderson /*
6440d99d37aSRichard Henderson  * Interval trees.
6450d99d37aSRichard Henderson  *
6460d99d37aSRichard Henderson  * Derived from lib/interval_tree.c and its dependencies,
6470d99d37aSRichard Henderson  * especially include/linux/interval_tree_generic.h.
6480d99d37aSRichard Henderson  */
6490d99d37aSRichard Henderson 
6500d99d37aSRichard Henderson #define rb_to_itree(N)  container_of(N, IntervalTreeNode, rb)
6510d99d37aSRichard Henderson 
interval_tree_compute_max(IntervalTreeNode * node,bool exit)6520d99d37aSRichard Henderson static bool interval_tree_compute_max(IntervalTreeNode *node, bool exit)
6530d99d37aSRichard Henderson {
6540d99d37aSRichard Henderson     IntervalTreeNode *child;
6550d99d37aSRichard Henderson     uint64_t max = node->last;
6560d99d37aSRichard Henderson 
6570d99d37aSRichard Henderson     if (node->rb.rb_left) {
6580d99d37aSRichard Henderson         child = rb_to_itree(node->rb.rb_left);
6590d99d37aSRichard Henderson         if (child->subtree_last > max) {
6600d99d37aSRichard Henderson             max = child->subtree_last;
6610d99d37aSRichard Henderson         }
6620d99d37aSRichard Henderson     }
6630d99d37aSRichard Henderson     if (node->rb.rb_right) {
6640d99d37aSRichard Henderson         child = rb_to_itree(node->rb.rb_right);
6650d99d37aSRichard Henderson         if (child->subtree_last > max) {
6660d99d37aSRichard Henderson             max = child->subtree_last;
6670d99d37aSRichard Henderson         }
6680d99d37aSRichard Henderson     }
6690d99d37aSRichard Henderson     if (exit && node->subtree_last == max) {
6700d99d37aSRichard Henderson         return true;
6710d99d37aSRichard Henderson     }
6720d99d37aSRichard Henderson     node->subtree_last = max;
6730d99d37aSRichard Henderson     return false;
6740d99d37aSRichard Henderson }
6750d99d37aSRichard Henderson 
interval_tree_propagate(RBNode * rb,RBNode * stop)6760d99d37aSRichard Henderson static void interval_tree_propagate(RBNode *rb, RBNode *stop)
6770d99d37aSRichard Henderson {
6780d99d37aSRichard Henderson     while (rb != stop) {
6790d99d37aSRichard Henderson         IntervalTreeNode *node = rb_to_itree(rb);
6800d99d37aSRichard Henderson         if (interval_tree_compute_max(node, true)) {
6810d99d37aSRichard Henderson             break;
6820d99d37aSRichard Henderson         }
6830d99d37aSRichard Henderson         rb = rb_parent(&node->rb);
6840d99d37aSRichard Henderson     }
6850d99d37aSRichard Henderson }
6860d99d37aSRichard Henderson 
interval_tree_copy(RBNode * rb_old,RBNode * rb_new)6870d99d37aSRichard Henderson static void interval_tree_copy(RBNode *rb_old, RBNode *rb_new)
6880d99d37aSRichard Henderson {
6890d99d37aSRichard Henderson     IntervalTreeNode *old = rb_to_itree(rb_old);
6900d99d37aSRichard Henderson     IntervalTreeNode *new = rb_to_itree(rb_new);
6910d99d37aSRichard Henderson 
6920d99d37aSRichard Henderson     new->subtree_last = old->subtree_last;
6930d99d37aSRichard Henderson }
6940d99d37aSRichard Henderson 
interval_tree_rotate(RBNode * rb_old,RBNode * rb_new)6950d99d37aSRichard Henderson static void interval_tree_rotate(RBNode *rb_old, RBNode *rb_new)
6960d99d37aSRichard Henderson {
6970d99d37aSRichard Henderson     IntervalTreeNode *old = rb_to_itree(rb_old);
6980d99d37aSRichard Henderson     IntervalTreeNode *new = rb_to_itree(rb_new);
6990d99d37aSRichard Henderson 
7000d99d37aSRichard Henderson     new->subtree_last = old->subtree_last;
7010d99d37aSRichard Henderson     interval_tree_compute_max(old, false);
7020d99d37aSRichard Henderson }
7030d99d37aSRichard Henderson 
7040d99d37aSRichard Henderson static const RBAugmentCallbacks interval_tree_augment = {
7050d99d37aSRichard Henderson     .propagate = interval_tree_propagate,
7060d99d37aSRichard Henderson     .copy = interval_tree_copy,
7070d99d37aSRichard Henderson     .rotate = interval_tree_rotate,
7080d99d37aSRichard Henderson };
7090d99d37aSRichard Henderson 
7100d99d37aSRichard Henderson /* Insert / remove interval nodes from the tree */
interval_tree_insert(IntervalTreeNode * node,IntervalTreeRoot * root)7110d99d37aSRichard Henderson void interval_tree_insert(IntervalTreeNode *node, IntervalTreeRoot *root)
7120d99d37aSRichard Henderson {
7130d99d37aSRichard Henderson     RBNode **link = &root->rb_root.rb_node, *rb_parent = NULL;
7140d99d37aSRichard Henderson     uint64_t start = node->start, last = node->last;
7150d99d37aSRichard Henderson     IntervalTreeNode *parent;
7160d99d37aSRichard Henderson     bool leftmost = true;
7170d99d37aSRichard Henderson 
7180d99d37aSRichard Henderson     while (*link) {
7190d99d37aSRichard Henderson         rb_parent = *link;
7200d99d37aSRichard Henderson         parent = rb_to_itree(rb_parent);
7210d99d37aSRichard Henderson 
7220d99d37aSRichard Henderson         if (parent->subtree_last < last) {
7230d99d37aSRichard Henderson             parent->subtree_last = last;
7240d99d37aSRichard Henderson         }
7250d99d37aSRichard Henderson         if (start < parent->start) {
7260d99d37aSRichard Henderson             link = &parent->rb.rb_left;
7270d99d37aSRichard Henderson         } else {
7280d99d37aSRichard Henderson             link = &parent->rb.rb_right;
7290d99d37aSRichard Henderson             leftmost = false;
7300d99d37aSRichard Henderson         }
7310d99d37aSRichard Henderson     }
7320d99d37aSRichard Henderson 
7330d99d37aSRichard Henderson     node->subtree_last = last;
7340d99d37aSRichard Henderson     rb_link_node(&node->rb, rb_parent, link);
7350d99d37aSRichard Henderson     rb_insert_augmented_cached(&node->rb, root, leftmost,
7360d99d37aSRichard Henderson                                &interval_tree_augment);
7370d99d37aSRichard Henderson }
7380d99d37aSRichard Henderson 
interval_tree_remove(IntervalTreeNode * node,IntervalTreeRoot * root)7390d99d37aSRichard Henderson void interval_tree_remove(IntervalTreeNode *node, IntervalTreeRoot *root)
7400d99d37aSRichard Henderson {
7410d99d37aSRichard Henderson     rb_erase_augmented_cached(&node->rb, root, &interval_tree_augment);
7420d99d37aSRichard Henderson }
7430d99d37aSRichard Henderson 
7440d99d37aSRichard Henderson /*
7450d99d37aSRichard Henderson  * Iterate over intervals intersecting [start;last]
7460d99d37aSRichard Henderson  *
7470d99d37aSRichard Henderson  * Note that a node's interval intersects [start;last] iff:
7480d99d37aSRichard Henderson  *   Cond1: node->start <= last
7490d99d37aSRichard Henderson  * and
7500d99d37aSRichard Henderson  *   Cond2: start <= node->last
7510d99d37aSRichard Henderson  */
7520d99d37aSRichard Henderson 
interval_tree_subtree_search(IntervalTreeNode * node,uint64_t start,uint64_t last)7530d99d37aSRichard Henderson static IntervalTreeNode *interval_tree_subtree_search(IntervalTreeNode *node,
7540d99d37aSRichard Henderson                                                       uint64_t start,
7550d99d37aSRichard Henderson                                                       uint64_t last)
7560d99d37aSRichard Henderson {
7570d99d37aSRichard Henderson     while (true) {
7580d99d37aSRichard Henderson         /*
7590d99d37aSRichard Henderson          * Loop invariant: start <= node->subtree_last
7600d99d37aSRichard Henderson          * (Cond2 is satisfied by one of the subtree nodes)
7610d99d37aSRichard Henderson          */
762055b86e0SRichard Henderson         RBNode *tmp = qatomic_read(&node->rb.rb_left);
763055b86e0SRichard Henderson         if (tmp) {
764055b86e0SRichard Henderson             IntervalTreeNode *left = rb_to_itree(tmp);
7650d99d37aSRichard Henderson 
7660d99d37aSRichard Henderson             if (start <= left->subtree_last) {
7670d99d37aSRichard Henderson                 /*
7680d99d37aSRichard Henderson                  * Some nodes in left subtree satisfy Cond2.
7690d99d37aSRichard Henderson                  * Iterate to find the leftmost such node N.
7700d99d37aSRichard Henderson                  * If it also satisfies Cond1, that's the
7710d99d37aSRichard Henderson                  * match we are looking for. Otherwise, there
7720d99d37aSRichard Henderson                  * is no matching interval as nodes to the
7730d99d37aSRichard Henderson                  * right of N can't satisfy Cond1 either.
7740d99d37aSRichard Henderson                  */
7750d99d37aSRichard Henderson                 node = left;
7760d99d37aSRichard Henderson                 continue;
7770d99d37aSRichard Henderson             }
7780d99d37aSRichard Henderson         }
7790d99d37aSRichard Henderson         if (node->start <= last) {         /* Cond1 */
7800d99d37aSRichard Henderson             if (start <= node->last) {     /* Cond2 */
7810d99d37aSRichard Henderson                 return node; /* node is leftmost match */
7820d99d37aSRichard Henderson             }
783055b86e0SRichard Henderson             tmp = qatomic_read(&node->rb.rb_right);
784055b86e0SRichard Henderson             if (tmp) {
785055b86e0SRichard Henderson                 node = rb_to_itree(tmp);
7860d99d37aSRichard Henderson                 if (start <= node->subtree_last) {
7870d99d37aSRichard Henderson                     continue;
7880d99d37aSRichard Henderson                 }
7890d99d37aSRichard Henderson             }
7900d99d37aSRichard Henderson         }
7910d99d37aSRichard Henderson         return NULL; /* no match */
7920d99d37aSRichard Henderson     }
7930d99d37aSRichard Henderson }
7940d99d37aSRichard Henderson 
interval_tree_iter_first(IntervalTreeRoot * root,uint64_t start,uint64_t last)7950d99d37aSRichard Henderson IntervalTreeNode *interval_tree_iter_first(IntervalTreeRoot *root,
7960d99d37aSRichard Henderson                                            uint64_t start, uint64_t last)
7970d99d37aSRichard Henderson {
7980d99d37aSRichard Henderson     IntervalTreeNode *node, *leftmost;
7990d99d37aSRichard Henderson 
80047d1e982SHelge Deller     if (!root || !root->rb_root.rb_node) {
8010d99d37aSRichard Henderson         return NULL;
8020d99d37aSRichard Henderson     }
8030d99d37aSRichard Henderson 
8040d99d37aSRichard Henderson     /*
8050d99d37aSRichard Henderson      * Fastpath range intersection/overlap between A: [a0, a1] and
8060d99d37aSRichard Henderson      * B: [b0, b1] is given by:
8070d99d37aSRichard Henderson      *
8080d99d37aSRichard Henderson      *         a0 <= b1 && b0 <= a1
8090d99d37aSRichard Henderson      *
8100d99d37aSRichard Henderson      *  ... where A holds the lock range and B holds the smallest
8110d99d37aSRichard Henderson      * 'start' and largest 'last' in the tree. For the later, we
8120d99d37aSRichard Henderson      * rely on the root node, which by augmented interval tree
8130d99d37aSRichard Henderson      * property, holds the largest value in its last-in-subtree.
8140d99d37aSRichard Henderson      * This allows mitigating some of the tree walk overhead for
8150d99d37aSRichard Henderson      * for non-intersecting ranges, maintained and consulted in O(1).
8160d99d37aSRichard Henderson      */
8170d99d37aSRichard Henderson     node = rb_to_itree(root->rb_root.rb_node);
8180d99d37aSRichard Henderson     if (node->subtree_last < start) {
8190d99d37aSRichard Henderson         return NULL;
8200d99d37aSRichard Henderson     }
8210d99d37aSRichard Henderson 
8220d99d37aSRichard Henderson     leftmost = rb_to_itree(root->rb_leftmost);
8230d99d37aSRichard Henderson     if (leftmost->start > last) {
8240d99d37aSRichard Henderson         return NULL;
8250d99d37aSRichard Henderson     }
8260d99d37aSRichard Henderson 
8270d99d37aSRichard Henderson     return interval_tree_subtree_search(node, start, last);
8280d99d37aSRichard Henderson }
8290d99d37aSRichard Henderson 
interval_tree_iter_next(IntervalTreeNode * node,uint64_t start,uint64_t last)8300d99d37aSRichard Henderson IntervalTreeNode *interval_tree_iter_next(IntervalTreeNode *node,
8310d99d37aSRichard Henderson                                           uint64_t start, uint64_t last)
8320d99d37aSRichard Henderson {
833055b86e0SRichard Henderson     RBNode *rb, *prev;
8340d99d37aSRichard Henderson 
835055b86e0SRichard Henderson     rb = qatomic_read(&node->rb.rb_right);
8360d99d37aSRichard Henderson     while (true) {
8370d99d37aSRichard Henderson         /*
8380d99d37aSRichard Henderson          * Loop invariants:
8390d99d37aSRichard Henderson          *   Cond1: node->start <= last
8400d99d37aSRichard Henderson          *   rb == node->rb.rb_right
8410d99d37aSRichard Henderson          *
8420d99d37aSRichard Henderson          * First, search right subtree if suitable
8430d99d37aSRichard Henderson          */
8440d99d37aSRichard Henderson         if (rb) {
8450d99d37aSRichard Henderson             IntervalTreeNode *right = rb_to_itree(rb);
8460d99d37aSRichard Henderson 
8470d99d37aSRichard Henderson             if (start <= right->subtree_last) {
8480d99d37aSRichard Henderson                 return interval_tree_subtree_search(right, start, last);
8490d99d37aSRichard Henderson             }
8500d99d37aSRichard Henderson         }
8510d99d37aSRichard Henderson 
8520d99d37aSRichard Henderson         /* Move up the tree until we come from a node's left child */
8530d99d37aSRichard Henderson         do {
8540d99d37aSRichard Henderson             rb = rb_parent(&node->rb);
8550d99d37aSRichard Henderson             if (!rb) {
8560d99d37aSRichard Henderson                 return NULL;
8570d99d37aSRichard Henderson             }
8580d99d37aSRichard Henderson             prev = &node->rb;
8590d99d37aSRichard Henderson             node = rb_to_itree(rb);
860055b86e0SRichard Henderson             rb = qatomic_read(&node->rb.rb_right);
8610d99d37aSRichard Henderson         } while (prev == rb);
8620d99d37aSRichard Henderson 
8630d99d37aSRichard Henderson         /* Check if the node intersects [start;last] */
8640d99d37aSRichard Henderson         if (last < node->start) {  /* !Cond1 */
8650d99d37aSRichard Henderson             return NULL;
8660d99d37aSRichard Henderson         }
8670d99d37aSRichard Henderson         if (start <= node->last) { /* Cond2 */
8680d99d37aSRichard Henderson             return node;
8690d99d37aSRichard Henderson         }
8700d99d37aSRichard Henderson     }
8710d99d37aSRichard Henderson }
8720d99d37aSRichard Henderson 
8730d99d37aSRichard Henderson /* Occasionally useful for calling from within the debugger. */
8740d99d37aSRichard Henderson #if 0
8750d99d37aSRichard Henderson static void debug_interval_tree_int(IntervalTreeNode *node,
8760d99d37aSRichard Henderson                                     const char *dir, int level)
8770d99d37aSRichard Henderson {
8780d99d37aSRichard Henderson     printf("%4d %*s %s [%" PRIu64 ",%" PRIu64 "] subtree_last:%" PRIu64 "\n",
8790d99d37aSRichard Henderson            level, level + 1, dir, rb_is_red(&node->rb) ? "r" : "b",
8800d99d37aSRichard Henderson            node->start, node->last, node->subtree_last);
8810d99d37aSRichard Henderson 
8820d99d37aSRichard Henderson     if (node->rb.rb_left) {
8830d99d37aSRichard Henderson         debug_interval_tree_int(rb_to_itree(node->rb.rb_left), "<", level + 1);
8840d99d37aSRichard Henderson     }
8850d99d37aSRichard Henderson     if (node->rb.rb_right) {
8860d99d37aSRichard Henderson         debug_interval_tree_int(rb_to_itree(node->rb.rb_right), ">", level + 1);
8870d99d37aSRichard Henderson     }
8880d99d37aSRichard Henderson }
8890d99d37aSRichard Henderson 
8900d99d37aSRichard Henderson void debug_interval_tree(IntervalTreeNode *node);
8910d99d37aSRichard Henderson void debug_interval_tree(IntervalTreeNode *node)
8920d99d37aSRichard Henderson {
8930d99d37aSRichard Henderson     if (node) {
8940d99d37aSRichard Henderson         debug_interval_tree_int(node, "*", 0);
8950d99d37aSRichard Henderson     } else {
8960d99d37aSRichard Henderson         printf("null\n");
8970d99d37aSRichard Henderson     }
8980d99d37aSRichard Henderson }
8990d99d37aSRichard Henderson #endif
900