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 RB_DECLARE_CALLBACKS_MAX(static, ITPREFIX ## _augment,                        \
  34                          ITSTRUCT, ITRB, ITTYPE, ITSUBTREE, ITLAST)           \
  35                                                                               \
  36 /* Insert / remove interval nodes from the tree */                            \
  37                                                                               \
  38 ITSTATIC void ITPREFIX ## _insert(ITSTRUCT *node,                             \
  39                                   struct rb_root_cached *root)                \
  40 {                                                                             \
  41         struct rb_node **link = &root->rb_root.rb_node, *rb_parent = NULL;    \
  42         ITTYPE start = ITSTART(node), last = ITLAST(node);                    \
  43         ITSTRUCT *parent;                                                     \
  44         bool leftmost = true;                                                 \
  45                                                                               \
  46         while (*link) {                                                       \
  47                 rb_parent = *link;                                            \
  48                 parent = rb_entry(rb_parent, ITSTRUCT, ITRB);                 \
  49                 if (parent->ITSUBTREE < last)                                 \
  50                         parent->ITSUBTREE = last;                             \
  51                 if (start < ITSTART(parent))                                  \
  52                         link = &parent->ITRB.rb_left;                         \
  53                 else {                                                        \
  54                         link = &parent->ITRB.rb_right;                        \
  55                         leftmost = false;                                     \
  56                 }                                                             \
  57         }                                                                     \
  58                                                                               \
  59         node->ITSUBTREE = last;                                               \
  60         rb_link_node(&node->ITRB, rb_parent, link);                           \
  61         rb_insert_augmented_cached(&node->ITRB, root,                         \
  62                                    leftmost, &ITPREFIX ## _augment);          \
  63 }                                                                             \
  64                                                                               \
  65 ITSTATIC void ITPREFIX ## _remove(ITSTRUCT *node,                             \
  66                                   struct rb_root_cached *root)                \
  67 {                                                                             \
  68         rb_erase_augmented_cached(&node->ITRB, root, &ITPREFIX ## _augment);  \
  69 }                                                                             \
  70                                                                               \
  71 /*                                                                            \
  72  * Iterate over intervals intersecting [start;last]                           \
  73  *                                                                            \
  74  * Note that a node's interval intersects [start;last] iff:                   \
  75  *   Cond1: ITSTART(node) <= last                                             \
  76  * and                                                                        \
  77  *   Cond2: start <= ITLAST(node)                                             \
  78  */                                                                           \
  79                                                                               \
  80 static ITSTRUCT *                                                             \
  81 ITPREFIX ## _subtree_search(ITSTRUCT *node, ITTYPE start, ITTYPE last)        \
  82 {                                                                             \
  83         while (true) {                                                        \
  84                 /*                                                            \
  85                  * Loop invariant: start <= node->ITSUBTREE                   \
  86                  * (Cond2 is satisfied by one of the subtree nodes)           \
  87                  */                                                           \
  88                 if (node->ITRB.rb_left) {                                     \
  89                         ITSTRUCT *left = rb_entry(node->ITRB.rb_left,         \
  90                                                   ITSTRUCT, ITRB);            \
  91                         if (start <= left->ITSUBTREE) {                       \
  92                                 /*                                            \
  93                                  * Some nodes in left subtree satisfy Cond2.  \
  94                                  * Iterate to find the leftmost such node N.  \
  95                                  * If it also satisfies Cond1, that's the     \
  96                                  * match we are looking for. Otherwise, there \
  97                                  * is no matching interval as nodes to the    \
  98                                  * right of N can't satisfy Cond1 either.     \
  99                                  */                                           \
 100                                 node = left;                                  \
 101                                 continue;                                     \
 102                         }                                                     \
 103                 }                                                             \
 104                 if (ITSTART(node) <= last) {            /* Cond1 */           \
 105                         if (start <= ITLAST(node))      /* Cond2 */           \
 106                                 return node;    /* node is leftmost match */  \
 107                         if (node->ITRB.rb_right) {                            \
 108                                 node = rb_entry(node->ITRB.rb_right,          \
 109                                                 ITSTRUCT, ITRB);              \
 110                                 if (start <= node->ITSUBTREE)                 \
 111                                         continue;                             \
 112                         }                                                     \
 113                 }                                                             \
 114                 return NULL;    /* No match */                                \
 115         }                                                                     \
 116 }                                                                             \
 117                                                                               \
 118 ITSTATIC ITSTRUCT *                                                           \
 119 ITPREFIX ## _iter_first(struct rb_root_cached *root,                          \
 120                         ITTYPE start, ITTYPE last)                            \
 121 {                                                                             \
 122         ITSTRUCT *node, *leftmost;                                            \
 123                                                                               \
 124         if (!root->rb_root.rb_node)                                           \
 125                 return NULL;                                                  \
 126                                                                               \
 127         /*                                                                    \
 128          * Fastpath range intersection/overlap between A: [a0, a1] and        \
 129          * B: [b0, b1] is given by:                                           \
 130          *                                                                    \
 131          *         a0 <= b1 && b0 <= a1                                       \
 132          *                                                                    \
 133          *  ... where A holds the lock range and B holds the smallest         \
 134          * 'start' and largest 'last' in the tree. For the later, we          \
 135          * rely on the root node, which by augmented interval tree            \
 136          * property, holds the largest value in its last-in-subtree.          \
 137          * This allows mitigating some of the tree walk overhead for          \
 138          * for non-intersecting ranges, maintained and consulted in O(1).     \
 139          */                                                                   \
 140         node = rb_entry(root->rb_root.rb_node, ITSTRUCT, ITRB);               \
 141         if (node->ITSUBTREE < start)                                          \
 142                 return NULL;                                                  \
 143                                                                               \
 144         leftmost = rb_entry(root->rb_leftmost, ITSTRUCT, ITRB);               \
 145         if (ITSTART(leftmost) > last)                                         \
 146                 return NULL;                                                  \
 147                                                                               \
 148         return ITPREFIX ## _subtree_search(node, start, last);                \
 149 }                                                                             \
 150                                                                               \
 151 ITSTATIC ITSTRUCT *                                                           \
 152 ITPREFIX ## _iter_next(ITSTRUCT *node, ITTYPE start, ITTYPE last)             \
 153 {                                                                             \
 154         struct rb_node *rb = node->ITRB.rb_right, *prev;                      \
 155                                                                               \
 156         while (true) {                                                        \
 157                 /*                                                            \
 158                  * Loop invariants:                                           \
 159                  *   Cond1: ITSTART(node) <= last                             \
 160                  *   rb == node->ITRB.rb_right                                \
 161                  *                                                            \
 162                  * First, search right subtree if suitable                    \
 163                  */                                                           \
 164                 if (rb) {                                                     \
 165                         ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB);       \
 166                         if (start <= right->ITSUBTREE)                        \
 167                                 return ITPREFIX ## _subtree_search(right,     \
 168                                                                 start, last); \
 169                 }                                                             \
 170                                                                               \
 171                 /* Move up the tree until we come from a node's left child */ \
 172                 do {                                                          \
 173                         rb = rb_parent(&node->ITRB);                          \
 174                         if (!rb)                                              \
 175                                 return NULL;                                  \
 176                         prev = &node->ITRB;                                   \
 177                         node = rb_entry(rb, ITSTRUCT, ITRB);                  \
 178                         rb = node->ITRB.rb_right;                             \
 179                 } while (prev == rb);                                         \
 180                                                                               \
 181                 /* Check if the node intersects [start;last] */               \
 182                 if (last < ITSTART(node))               /* !Cond1 */          \
 183                         return NULL;                                          \
 184                 else if (start <= ITLAST(node))         /* Cond2 */           \
 185                         return node;                                          \
 186         }                                                                     \
 187 }