include/linux/interval_tree_tmpl.h

   1 /*
   2   Interval Trees
   3   (C) 2012  Michel Lespinasse <walken@google.com>
   4
   5   This program is free software; you can redistribute it and/or modify
   6   it under the terms of the GNU General Public License as published by
   7   the Free Software Foundation; either version 2 of the License, or
   8   (at your option) any later version.
   9
  10   This program is distributed in the hope that it will be useful,
  11   but WITHOUT ANY WARRANTY; without even the implied warranty of
  12   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13   GNU General Public License for more details.
  14
  15   You should have received a copy of the GNU General Public License
  16   along with this program; if not, write to the Free Software
  17   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  18
  19   include/linux/interval_tree_tmpl.h
  20 */
  21
  22 #include <linux/rbtree_augmented.h>
  23
  24 /*
  25  * Template for implementing interval trees
  26  *
  27  * ITSTRUCT:   struct type of the interval tree nodes
  28  * ITRB:       name of struct rb_node field within ITSTRUCT
  29  * ITTYPE:     type of the interval endpoints
  30  * ITSUBTREE:  name of ITTYPE field within ITSTRUCT holding last-in-subtree
  31  * ITSTART(n): start endpoint of ITSTRUCT node n
  32  * ITLAST(n):  last endpoing of ITSTRUCT node n
  33  * ITSTATIC:   'static' or empty
  34  * ITPREFIX:   prefix to use for the inline tree definitions
  35  */
  36
  37 /* IT(name) -> ITPREFIX_name */
  38 #define _ITNAME(prefix, name) prefix ## _ ## name
  39 #define ITNAME(prefix, name) _ITNAME(prefix, name)
  40 #define IT(name) ITNAME(ITPREFIX, name)
  41
  42 /* Callbacks for augmented rbtree insert and remove */
  43
  44 static inline ITTYPE IT(compute_subtree_last)(ITSTRUCT *node)
  45 {
  46         ITTYPE max = ITLAST(node), subtree_last;
  47         if (node->ITRB.rb_left) {
  48                 subtree_last = rb_entry(node->ITRB.rb_left,
  49                                         ITSTRUCT, ITRB)->ITSUBTREE;
  50                 if (max < subtree_last)
  51                         max = subtree_last;
  52         }
  53         if (node->ITRB.rb_right) {
  54                 subtree_last = rb_entry(node->ITRB.rb_right,
  55                                         ITSTRUCT, ITRB)->ITSUBTREE;
  56                 if (max < subtree_last)
  57                         max = subtree_last;
  58         }
  59         return max;
  60 }
  61
  62 static inline void
  63 IT(augment_propagate)(struct rb_node *rb, struct rb_node *stop)
  64 {
  65         while (rb != stop) {
  66                 ITSTRUCT *node = rb_entry(rb, ITSTRUCT, ITRB);
  67                 ITTYPE subtree_last = IT(compute_subtree_last)(node);
  68                 if (node->ITSUBTREE == subtree_last)
  69                         break;
  70                 node->ITSUBTREE = subtree_last;
  71                 rb = rb_parent(&node->ITRB);
  72         }
  73 }
  74
  75 static inline void
  76 IT(augment_copy)(struct rb_node *rb_old, struct rb_node *rb_new)
  77 {
  78         ITSTRUCT *old = rb_entry(rb_old, ITSTRUCT, ITRB);
  79         ITSTRUCT *new = rb_entry(rb_new, ITSTRUCT, ITRB);
  80
  81         new->ITSUBTREE = old->ITSUBTREE;
  82 }
  83
  84 static void IT(augment_rotate)(struct rb_node *rb_old, struct rb_node *rb_new)
  85 {
  86         ITSTRUCT *old = rb_entry(rb_old, ITSTRUCT, ITRB);
  87         ITSTRUCT *new = rb_entry(rb_new, ITSTRUCT, ITRB);
  88
  89         new->ITSUBTREE = old->ITSUBTREE;
  90         old->ITSUBTREE = IT(compute_subtree_last)(old);
  91 }
  92
  93 static const struct rb_augment_callbacks IT(augment_callbacks) = {
  94         IT(augment_propagate), IT(augment_copy), IT(augment_rotate)
  95 };
  96
  97 /* Insert / remove interval nodes from the tree */
  98
  99 ITSTATIC void IT(insert)(ITSTRUCT *node, struct rb_root *root)
 100 {
 101         struct rb_node **link = &root->rb_node, *rb_parent = NULL;
 102         ITTYPE start = ITSTART(node), last = ITLAST(node);
 103         ITSTRUCT *parent;
 104
 105         while (*link) {
 106                 rb_parent = *link;
 107                 parent = rb_entry(rb_parent, ITSTRUCT, ITRB);
 108                 if (parent->ITSUBTREE < last)
 109                         parent->ITSUBTREE = last;
 110                 if (start < ITSTART(parent))
 111                         link = &parent->ITRB.rb_left;
 112                 else
 113                         link = &parent->ITRB.rb_right;
 114         }
 115
 116         node->ITSUBTREE = last;
 117         rb_link_node(&node->ITRB, rb_parent, link);
 118         rb_insert_augmented(&node->ITRB, root, &IT(augment_callbacks));
 119 }
 120
 121 ITSTATIC void IT(remove)(ITSTRUCT *node, struct rb_root *root)
 122 {
 123         rb_erase_augmented(&node->ITRB, root, &IT(augment_callbacks));
 124 }
 125
 126 /*
 127  * Iterate over intervals intersecting [start;last]
 128  *
 129  * Note that a node's interval intersects [start;last] iff:
 130  *   Cond1: ITSTART(node) <= last
 131  * and
 132  *   Cond2: start <= ITLAST(node)
 133  */
 134
 135 static ITSTRUCT *IT(subtree_search)(ITSTRUCT *node, ITTYPE start, ITTYPE last)
 136 {
 137         while (true) {
 138                 /*
 139                  * Loop invariant: start <= node->ITSUBTREE
 140                  * (Cond2 is satisfied by one of the subtree nodes)
 141                  */
 142                 if (node->ITRB.rb_left) {
 143                         ITSTRUCT *left = rb_entry(node->ITRB.rb_left,
 144                                                   ITSTRUCT, ITRB);
 145                         if (start <= left->ITSUBTREE) {
 146                                 /*
 147                                  * Some nodes in left subtree satisfy Cond2.
 148                                  * Iterate to find the leftmost such node N.
 149                                  * If it also satisfies Cond1, that's the match
 150                                  * we are looking for. Otherwise, there is no
 151                                  * matching interval as nodes to the right of N
 152                                  * can't satisfy Cond1 either.
 153                                  */
 154                                 node = left;
 155                                 continue;
 156                         }
 157                 }
 158                 if (ITSTART(node) <= last) {            /* Cond1 */
 159                         if (start <= ITLAST(node))      /* Cond2 */
 160                                 return node;    /* node is leftmost match */
 161                         if (node->ITRB.rb_right) {
 162                                 node = rb_entry(node->ITRB.rb_right,
 163                                                 ITSTRUCT, ITRB);
 164                                 if (start <= node->ITSUBTREE)
 165                                         continue;
 166                         }
 167                 }
 168                 return NULL;    /* No match */
 169         }
 170 }
 171
 172 ITSTATIC ITSTRUCT *IT(iter_first)(struct rb_root *root,
 173                                   ITTYPE start, ITTYPE last)
 174 {
 175         ITSTRUCT *node;
 176
 177         if (!root->rb_node)
 178                 return NULL;
 179         node = rb_entry(root->rb_node, ITSTRUCT, ITRB);
 180         if (node->ITSUBTREE < start)
 181                 return NULL;
 182         return IT(subtree_search)(node, start, last);
 183 }
 184
 185 ITSTATIC ITSTRUCT *IT(iter_next)(ITSTRUCT *node, ITTYPE start, ITTYPE last)
 186 {
 187         struct rb_node *rb = node->ITRB.rb_right, *prev;
 188
 189         while (true) {
 190                 /*
 191                  * Loop invariants:
 192                  *   Cond1: ITSTART(node) <= last
 193                  *   rb == node->ITRB.rb_right
 194                  *
 195                  * First, search right subtree if suitable
 196                  */
 197                 if (rb) {
 198                         ITSTRUCT *right = rb_entry(rb, ITSTRUCT, ITRB);
 199                         if (start <= right->ITSUBTREE)
 200                                 return IT(subtree_search)(right, start, last);
 201                 }
 202
 203                 /* Move up the tree until we come from a node's left child */
 204                 do {
 205                         rb = rb_parent(&node->ITRB);
 206                         if (!rb)
 207                                 return NULL;
 208                         prev = &node->ITRB;
 209                         node = rb_entry(rb, ITSTRUCT, ITRB);
 210                         rb = node->ITRB.rb_right;
 211                 } while (prev == rb);
 212
 213                 /* Check if the node intersects [start;last] */
 214                 if (last < ITSTART(node))               /* !Cond1 */
 215                         return NULL;
 216                 else if (start <= ITLAST(node))         /* Cond2 */
 217                         return node;
 218         }
 219 }