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