LXR linux/include/linux/interval_tree

   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                                                                              \
  33RB_DECLARE_CALLBACKS_MAX(static, ITPREFIX ## _augment,                        \
  34                         ITSTRUCT, ITRB, ITTYPE, ITSUBTREE, ITLAST)           \
  35                                                                              \
  36/* Insert / remove interval nodes from the tree */                            \
  37                                                                              \
  38ITSTATIC 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                                                                              \
  65ITSTATIC 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                                                                              \
  80static ITSTRUCT *                                                             \
  81ITPREFIX ## _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                                                                              \
 118ITSTATIC ITSTRUCT *                                                           \
 119ITPREFIX ## _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                                                                              \
 151ITSTATIC ITSTRUCT *                                                           \
 152ITPREFIX ## _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}
 188