include/linux/interval_tree_generic.h

   1 /* SPDX-License-Identifier: GPL-2.0-or-later */
   2 /*
   3   Interval Trees
   4   (C) 2012  Michel Lespinasse <walken@google.com>
   5
   6
   7   include/linux/interval_tree_generic.h
   8 */
   9
  10 #include <linux/rbtree_augmented.h>
  11
  12 /*
  13  * Template for implementing interval trees
  14  *
  15  * ITSTRUCT:   struct type of the interval tree nodes
  16  * ITRB:       name of struct rb_node field within ITSTRUCT
  17  * ITTYPE:     type of the interval endpoints
  18  * ITSUBTREE:  name of ITTYPE field within ITSTRUCT holding last-in-subtree
  19  * ITSTART(n): start endpoint of ITSTRUCT node n
  20  * ITLAST(n):  last endpoint of ITSTRUCT node n
  21  * ITSTATIC:   'static' or empty
  22  * ITPREFIX:   prefix to use for the inline tree definitions
  23  *
  24  * Note - before using this, please consider if generic version
  25  * (interval_tree.h) would work for you...
  26  */
  27
  28 #define INTERVAL_TREE_DEFINE(ITSTRUCT, ITRB, ITTYPE, ITSUBTREE,               \
  29                              ITSTART, ITLAST, ITSTATIC, ITPREFIX)             \
  30                                                                               \
  31 /* Callbacks for augmented rbtree insert and remove */                        \
  32                                                                               \
  33 static inline ITTYPE ITPREFIX ## _compute_subtree_last(ITSTRUCT *node)        \
  34 {                                                                             \
  35         ITTYPE max = ITLAST(node), subtree_last;                              \
  36         if (node->ITRB.rb_left) {                                             \
  37                 subtree_last = rb_entry(node->ITRB.rb_left,                   \
  38                                         ITSTRUCT, ITRB)->ITSUBTREE;           \
  39                 if (max < subtree_last)                                       \
  40                         max = subtree_last;                                   \
  41         }                                                                     \
  42         if (node->ITRB.rb_right) {                                            \
  43                 subtree_last = rb_entry(node->ITRB.rb_right,                  \
  44                                         ITSTRUCT, ITRB)->ITSUBTREE;           \
  45                 if (max < subtree_last)                                       \
  46                         max = subtree_last;                                   \
  47         }                                                                     \
  48         return max;                                                           \
  49 }                                                                             \
  50                                                                               \
  51 RB_DECLARE_CALLBACKS(static, ITPREFIX ## _augment, ITSTRUCT, ITRB,            \
  52                      ITTYPE, ITSUBTREE, ITPREFIX ## _compute_subtree_last)    \
  53                                                                               \
  54 /* Insert / remove interval nodes from the tree */                            \
  55                                                                               \
  56 ITSTATIC void ITPREFIX ## _insert(ITSTRUCT *node,                             \
  57                                   struct rb_root_cached *root)                \
  58 {                                                                             \
  59         struct rb_node **link = &root->rb_root.rb_node, *rb_parent = NULL;    \
  60         ITTYPE start = ITSTART(node), last = ITLAST(node);                    \
  61         ITSTRUCT *parent;                                                     \
  62         bool leftmost = true;                                                 \
  63                                                                               \
  64         while (*link) {                                                       \
  65                 rb_parent = *link;                                            \
  66                 parent = rb_entry(rb_parent, ITSTRUCT, ITRB);                 \
  67                 if (parent->ITSUBTREE < last)                                 \
  68                         parent->ITSUBTREE = last;                             \
  69                 if (start < ITSTART(parent))                                  \
  70                         link = &parent->ITRB.rb_left;                         \
  71                 else {                                                        \
  72                         link = &parent->ITRB.rb_right;                        \
  73                         leftmost = false;                                     \
  74                 }                                                             \
  75         }                                                                     \
  76                                                                               \
  77         node->ITSUBTREE = last;                                               \
  78         rb_link_node(&node->ITRB, rb_parent, link);                           \
  79         rb_insert_augmented_cached(&node->ITRB, root,                         \
  80                                    leftmost, &ITPREFIX ## _augment);          \
  81 }                                                                             \
  82                                                                               \
  83 ITSTATIC void ITPREFIX ## _remove(ITSTRUCT *node,                             \
  84                                   struct rb_root_cached *root)                \
  85 {                                                                             \
  86         rb_erase_augmented_cached(&node->ITRB, root, &ITPREFIX ## _augment);  \
  87 }                                                                             \
  88                                                                               \
  89 /*                                                                            \
  90  * Iterate over intervals intersecting [start;last]                           \
  91  *                                                                            \
  92  * Note that a node's interval intersects [start;last] iff:                   \
  93  *   Cond1: ITSTART(node) <= last                                             \
  94  * and                                                                        \
  95  *   Cond2: start <= ITLAST(node)                                             \
  96  */                                                                           \
  97                                                                               \
  98 static ITSTRUCT *                                                             \
  99 ITPREFIX ## _subtree_search(ITSTRUCT *node, ITTYPE start, ITTYPE last)        \
 100 {                                                                             \
 101         while (true) {                                                        \
 102                 /*                                                            \
 103                  * Loop invariant: start <= node->ITSUBTREE                   \
 104                  * (Cond2 is satisfied by one of the subtree nodes)           \
 105                  */                                                           \
 106                 if (node->ITRB.rb_left) {                                     \
 107                         ITSTRUCT *left = rb_entry(node->ITRB.rb_left,         \
 108                                                   ITSTRUCT, ITRB);            \
 109                         if (start <= left->ITSUBTREE) {                       \
 110                                 /*                                            \
 111                                  * Some nodes in left subtree satisfy Cond2.  \
 112                                  * Iterate to find the leftmost such node N.  \
 113                                  * If it also satisfies Cond1, that's the     \
 114                                  * match we are looking for. Otherwise, there \
 115                                  * is no matching interval as nodes to the    \
 116                                  * right of N can't satisfy Cond1 either.     \
 117                                  */                                           \
 118                                 node = left;                                  \
 119                                 continue;                                     \
 120                         }                                                     \
 121                 }                                                             \
 122                 if (ITSTART(node) <= last) {            /* Cond1 */           \
 123                         if (start <= ITLAST(node))      /* Cond2 */           \
 124                                 return node;    /* node is leftmost match */  \
 125                         if (node->ITRB.rb_right) {                            \
 126                                 node = rb_entry(node->ITRB.rb_right,          \
 127                                                 ITSTRUCT, ITRB);              \
 128                                 if (start <= node->ITSUBTREE)                 \
 129                                         continue;                             \
 130                         }                                                     \
 131                 }                                                             \
 132                 return NULL;    /* No match */                                \
 133         }                                                                     \
 134 }                                                                             \
 135                                                                               \
 136 ITSTATIC ITSTRUCT *                                                           \
 137 ITPREFIX ## _iter_first(struct rb_root_cached *root,                          \
 138                         ITTYPE start, ITTYPE last)                            \
 139 {                                                                             \
 140         ITSTRUCT *node, *leftmost;                                            \
 141                                                                               \
 142         if (!root->rb_root.rb_node)                                           \
 143                 return NULL;                                                  \
 144                                                                               \
 145         /*                                                                    \
 146          * Fastpath range intersection/overlap between A: [a0, a1] and        \
 147          * B: [b0, b1] is given by:                                           \
 148          *                                                                    \
 149          *         a0 <= b1 && b0 <= a1                                       \
 150          *                                                                    \
 151          *  ... where A holds the lock range and B holds the smallest         \
 152          * 'start' and largest 'last' in the tree. For the later, we          \
 153          * rely on the root node, which by augmented interval tree            \
 154          * property, holds the largest value in its last-in-subtree.          \
 155          * This allows mitigating some of the tree walk overhead for          \
 156          * for non-intersecting ranges, maintained and consulted in O(1).     \
 157          */                                                                   \
 158         node = rb_entry(root->rb_root.rb_node, ITSTRUCT, ITRB);               \
 159         if (node->ITSUBTREE < start)                                          \
 160                 return NULL;                                                  \
 161                                                                               \
 162         leftmost = rb_entry(root->rb_leftmost, ITSTRUCT, ITRB);               \
 163         if (ITSTART(leftmost) > last)                                         \
 164                 return NULL;                                                  \
 165                                                                               \
 166         return ITPREFIX ## _subtree_search(node, start, last);                \
 167 }                                                                             \
 168                                                                               \
 169 ITSTATIC ITSTRUCT *                                                           \
 170 ITPREFIX ## _iter_next(ITSTRUCT *node, ITTYPE start, ITTYPE last)             \
 171 {                                                                             \
 172         struct rb_node *rb = node->ITRB.rb_right, *prev;                      \
 173                                                                               \
 174         while (true) {                                                        \
 175                 /*                                                            \
 176                  * Loop invariants:                                           \
 177                  *   Cond1: ITSTART(node) <= last                             \
 178                  *   rb == node->ITRB.rb_right                                \
 179                  *                                                            \
 180                  * First, search right subtree if suitable                    \
 181                  */                                                           \
 182                 if (rb) {                                                     \
 183                         ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB);       \
 184                         if (start <= right->ITSUBTREE)                        \
 185                                 return ITPREFIX ## _subtree_search(right,     \
 186                                                                 start, last); \
 187                 }                                                             \
 188                                                                               \
 189                 /* Move up the tree until we come from a node's left child */ \
 190                 do {                                                          \
 191                         rb = rb_parent(&node->ITRB);                          \
 192                         if (!rb)                                              \
 193                                 return NULL;                                  \
 194                         prev = &node->ITRB;                                   \
 195                         node = rb_entry(rb, ITSTRUCT, ITRB);                  \
 196                         rb = node->ITRB.rb_right;                             \
 197                 } while (prev == rb);                                         \
 198                                                                               \
 199                 /* Check if the node intersects [start;last] */               \
 200                 if (last < ITSTART(node))               /* !Cond1 */          \
 201                         return NULL;                                          \
 202                 else if (start <= ITLAST(node))         /* Cond2 */           \
 203                         return node;                                          \
 204         }                                                                     \
 205 }