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