include/linux/interval_tree_generic.h

0001 /* SPDX-License-Identifier: GPL-2.0-or-later */
0002 /*
0003   Interval Trees
0004   (C) 2012  Michel Lespinasse <walken@google.com>
0005
0006
0007   include/linux/interval_tree_generic.h
0008 */
0009
0010 #include <linux/rbtree_augmented.h>
0011
0012 /*
0013  * Template for implementing interval trees
0014  *
0015  * ITSTRUCT:   struct type of the interval tree nodes
0016  * ITRB:       name of struct rb_node field within ITSTRUCT
0017  * ITTYPE:     type of the interval endpoints
0018  * ITSUBTREE:  name of ITTYPE field within ITSTRUCT holding last-in-subtree
0019  * ITSTART(n): start endpoint of ITSTRUCT node n
0020  * ITLAST(n):  last endpoint of ITSTRUCT node n
0021  * ITSTATIC:   'static' or empty
0022  * ITPREFIX:   prefix to use for the inline tree definitions
0023  *
0024  * Note - before using this, please consider if generic version
0025  * (interval_tree.h) would work for you...
0026  */
0027
0028 #define INTERVAL_TREE_DEFINE(ITSTRUCT, ITRB, ITTYPE, ITSUBTREE,           \
0029                  ITSTART, ITLAST, ITSTATIC, ITPREFIX)         \
0030                                           \
0031 /* Callbacks for augmented rbtree insert and remove */                \
0032                                           \
0033 RB_DECLARE_CALLBACKS_MAX(static, ITPREFIX ## _augment,                \
0034              ITSTRUCT, ITRB, ITTYPE, ITSUBTREE, ITLAST)       \
0035                                           \
0036 /* Insert / remove interval nodes from the tree */                \
0037                                           \
0038 ITSTATIC void ITPREFIX ## _insert(ITSTRUCT *node,                 \
0039                   struct rb_root_cached *root)            \
0040 {                                         \
0041     struct rb_node **link = &root->rb_root.rb_node, *rb_parent = NULL;    \
0042     ITTYPE start = ITSTART(node), last = ITLAST(node);            \
0043     ITSTRUCT *parent;                             \
0044     bool leftmost = true;                             \
0045                                           \
0046     while (*link) {                               \
0047         rb_parent = *link;                        \
0048         parent = rb_entry(rb_parent, ITSTRUCT, ITRB);             \
0049         if (parent->ITSUBTREE < last)                     \
0050             parent->ITSUBTREE = last;                 \
0051         if (start < ITSTART(parent))                      \
0052             link = &parent->ITRB.rb_left;                 \
0053         else {                                \
0054             link = &parent->ITRB.rb_right;                \
0055             leftmost = false;                     \
0056         }                                 \
0057     }                                     \
0058                                           \
0059     node->ITSUBTREE = last;                           \
0060     rb_link_node(&node->ITRB, rb_parent, link);               \
0061     rb_insert_augmented_cached(&node->ITRB, root,                 \
0062                    leftmost, &ITPREFIX ## _augment);          \
0063 }                                         \
0064                                           \
0065 ITSTATIC void ITPREFIX ## _remove(ITSTRUCT *node,                 \
0066                   struct rb_root_cached *root)            \
0067 {                                         \
0068     rb_erase_augmented_cached(&node->ITRB, root, &ITPREFIX ## _augment);  \
0069 }                                         \
0070                                           \
0071 /*                                        \
0072  * Iterate over intervals intersecting [start;last]               \
0073  *                                        \
0074  * Note that a node's interval intersects [start;last] iff:           \
0075  *   Cond1: ITSTART(node) <= last                         \
0076  * and                                        \
0077  *   Cond2: start <= ITLAST(node)                         \
0078  */                                       \
0079                                           \
0080 static ITSTRUCT *                                 \
0081 ITPREFIX ## _subtree_search(ITSTRUCT *node, ITTYPE start, ITTYPE last)        \
0082 {                                         \
0083     while (true) {                                \
0084         /*                                \
0085          * Loop invariant: start <= node->ITSUBTREE           \
0086          * (Cond2 is satisfied by one of the subtree nodes)       \
0087          */                               \
0088         if (node->ITRB.rb_left) {                     \
0089             ITSTRUCT *left = rb_entry(node->ITRB.rb_left,         \
0090                           ITSTRUCT, ITRB);        \
0091             if (start <= left->ITSUBTREE) {               \
0092                 /*                        \
0093                  * Some nodes in left subtree satisfy Cond2.  \
0094                  * Iterate to find the leftmost such node N.  \
0095                  * If it also satisfies Cond1, that's the     \
0096                  * match we are looking for. Otherwise, there \
0097                  * is no matching interval as nodes to the    \
0098                  * right of N can't satisfy Cond1 either.     \
0099                  */                       \
0100                 node = left;                      \
0101                 continue;                     \
0102             }                             \
0103         }                                 \
0104         if (ITSTART(node) <= last) {        /* Cond1 */       \
0105             if (start <= ITLAST(node))  /* Cond2 */       \
0106                 return node;    /* node is leftmost match */  \
0107             if (node->ITRB.rb_right) {                \
0108                 node = rb_entry(node->ITRB.rb_right,          \
0109                         ITSTRUCT, ITRB);          \
0110                 if (start <= node->ITSUBTREE)             \
0111                     continue;                 \
0112             }                             \
0113         }                                 \
0114         return NULL;    /* No match */                    \
0115     }                                     \
0116 }                                         \
0117                                           \
0118 ITSTATIC ITSTRUCT *                               \
0119 ITPREFIX ## _iter_first(struct rb_root_cached *root,                  \
0120             ITTYPE start, ITTYPE last)                \
0121 {                                         \
0122     ITSTRUCT *node, *leftmost;                        \
0123                                           \
0124     if (!root->rb_root.rb_node)                       \
0125         return NULL;                              \
0126                                           \
0127     /*                                    \
0128      * Fastpath range intersection/overlap between A: [a0, a1] and        \
0129      * B: [b0, b1] is given by:                       \
0130      *                                    \
0131      *         a0 <= b1 && b0 <= a1                       \
0132      *                                    \
0133      *  ... where A holds the lock range and B holds the smallest         \
0134      * 'start' and largest 'last' in the tree. For the later, we          \
0135      * rely on the root node, which by augmented interval tree        \
0136      * property, holds the largest value in its last-in-subtree.          \
0137      * This allows mitigating some of the tree walk overhead for          \
0138      * for non-intersecting ranges, maintained and consulted in O(1).     \
0139      */                                   \
0140     node = rb_entry(root->rb_root.rb_node, ITSTRUCT, ITRB);           \
0141     if (node->ITSUBTREE < start)                          \
0142         return NULL;                              \
0143                                           \
0144     leftmost = rb_entry(root->rb_leftmost, ITSTRUCT, ITRB);           \
0145     if (ITSTART(leftmost) > last)                         \
0146         return NULL;                              \
0147                                           \
0148     return ITPREFIX ## _subtree_search(node, start, last);            \
0149 }                                         \
0150                                           \
0151 ITSTATIC ITSTRUCT *                               \
0152 ITPREFIX ## _iter_next(ITSTRUCT *node, ITTYPE start, ITTYPE last)         \
0153 {                                         \
0154     struct rb_node *rb = node->ITRB.rb_right, *prev;              \
0155                                           \
0156     while (true) {                                \
0157         /*                                \
0158          * Loop invariants:                       \
0159          *   Cond1: ITSTART(node) <= last                 \
0160          *   rb == node->ITRB.rb_right                    \
0161          *                                \
0162          * First, search right subtree if suitable            \
0163          */                               \
0164         if (rb) {                             \
0165             ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB);       \
0166             if (start <= right->ITSUBTREE)                \
0167                 return ITPREFIX ## _subtree_search(right,     \
0168                                 start, last); \
0169         }                                 \
0170                                           \
0171         /* Move up the tree until we come from a node's left child */ \
0172         do {                                  \
0173             rb = rb_parent(&node->ITRB);                  \
0174             if (!rb)                          \
0175                 return NULL;                      \
0176             prev = &node->ITRB;                   \
0177             node = rb_entry(rb, ITSTRUCT, ITRB);              \
0178             rb = node->ITRB.rb_right;                 \
0179         } while (prev == rb);                         \
0180                                           \
0181         /* Check if the node intersects [start;last] */           \
0182         if (last < ITSTART(node))       /* !Cond1 */          \
0183             return NULL;                          \
0184         else if (start <= ITLAST(node))     /* Cond2 */       \
0185             return node;                          \
0186     }                                     \
0187 }