tools/include/linux/rbtree_augmented.h

   1 /*
   2   Red Black Trees
   3   (C) 1999  Andrea Arcangeli <[email protected]>
   4   (C) 2002  David Woodhouse <[email protected]>
   5   (C) 2012  Michel Lespinasse <[email protected]>
   6
   7   This program is free software; you can redistribute it and/or modify
   8   it under the terms of the GNU General Public License as published by
   9   the Free Software Foundation; either version 2 of the License, or
  10   (at your option) any later version.
  11
  12   This program is distributed in the hope that it will be useful,
  13   but WITHOUT ANY WARRANTY; without even the implied warranty of
  14   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15   GNU General Public License for more details.
  16
  17   You should have received a copy of the GNU General Public License
  18   along with this program; if not, write to the Free Software
  19   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  20
  21   tools/linux/include/linux/rbtree_augmented.h
  22
  23   Copied from:
  24   linux/include/linux/rbtree_augmented.h
  25 */
  26
  27 #ifndef _TOOLS_LINUX_RBTREE_AUGMENTED_H
  28 #define _TOOLS_LINUX_RBTREE_AUGMENTED_H
  29
  30 #include <linux/compiler.h>
  31 #include <linux/rbtree.h>
  32
  33 /*
  34  * Please note - only struct rb_augment_callbacks and the prototypes for
  35  * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
  36  * The rest are implementation details you are not expected to depend on.
  37  *
  38  * See Documentation/rbtree.txt for documentation and samples.
  39  */
  40
  41 struct rb_augment_callbacks {
  42         void (*propagate)(struct rb_node *node, struct rb_node *stop);
  43         void (*copy)(struct rb_node *old, struct rb_node *new);
  44         void (*rotate)(struct rb_node *old, struct rb_node *new);
  45 };
  46
  47 extern void __rb_insert_augmented(struct rb_node *node,
  48                                   struct rb_root *root,
  49                                   bool newleft, struct rb_node **leftmost,
  50         void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
  51 /*
  52  * Fixup the rbtree and update the augmented information when rebalancing.
  53  *
  54  * On insertion, the user must update the augmented information on the path
  55  * leading to the inserted node, then call rb_link_node() as usual and
  56  * rb_augment_inserted() instead of the usual rb_insert_color() call.
  57  * If rb_augment_inserted() rebalances the rbtree, it will callback into
  58  * a user provided function to update the augmented information on the
  59  * affected subtrees.
  60  */
  61 static inline void
  62 rb_insert_augmented(struct rb_node *node, struct rb_root *root,
  63                     const struct rb_augment_callbacks *augment)
  64 {
  65         __rb_insert_augmented(node, root, false, NULL, augment->rotate);
  66 }
  67
  68 static inline void
  69 rb_insert_augmented_cached(struct rb_node *node,
  70                            struct rb_root_cached *root, bool newleft,
  71                            const struct rb_augment_callbacks *augment)
  72 {
  73         __rb_insert_augmented(node, &root->rb_root,
  74                               newleft, &root->rb_leftmost, augment->rotate);
  75 }
  76
  77 #define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield,       \
  78                              rbtype, rbaugmented, rbcompute)            \
  79 static inline void                                                      \
  80 rbname ## _propagate(struct rb_node *rb, struct rb_node *stop)          \
  81 {                                                                       \
  82         while (rb != stop) {                                            \
  83                 rbstruct *node = rb_entry(rb, rbstruct, rbfield);       \
  84                 rbtype augmented = rbcompute(node);                     \
  85                 if (node->rbaugmented == augmented)                     \
  86                         break;                                          \
  87                 node->rbaugmented = augmented;                          \
  88                 rb = rb_parent(&node->rbfield);                         \
  89         }                                                               \
  90 }                                                                       \
  91 static inline void                                                      \
  92 rbname ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)         \
  93 {                                                                       \
  94         rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);            \
  95         rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);            \
  96         new->rbaugmented = old->rbaugmented;                            \
  97 }                                                                       \
  98 static void                                                             \
  99 rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)       \
 100 {                                                                       \
 101         rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);            \
 102         rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);            \
 103         new->rbaugmented = old->rbaugmented;                            \
 104         old->rbaugmented = rbcompute(old);                              \
 105 }                                                                       \
 106 rbstatic const struct rb_augment_callbacks rbname = {                   \
 107         .propagate = rbname ## _propagate,                              \
 108         .copy = rbname ## _copy,                                        \
 109         .rotate = rbname ## _rotate                                     \
 110 };
 111
 112
 113 #define RB_RED          0
 114 #define RB_BLACK        1
 115
 116 #define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))
 117
 118 #define __rb_color(pc)     ((pc) & 1)
 119 #define __rb_is_black(pc)  __rb_color(pc)
 120 #define __rb_is_red(pc)    (!__rb_color(pc))
 121 #define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
 122 #define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
 123 #define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)
 124
 125 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
 126 {
 127         rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
 128 }
 129
 130 static inline void rb_set_parent_color(struct rb_node *rb,
 131                                        struct rb_node *p, int color)
 132 {
 133         rb->__rb_parent_color = (unsigned long)p | color;
 134 }
 135
 136 static inline void
 137 __rb_change_child(struct rb_node *old, struct rb_node *new,
 138                   struct rb_node *parent, struct rb_root *root)
 139 {
 140         if (parent) {
 141                 if (parent->rb_left == old)
 142                         WRITE_ONCE(parent->rb_left, new);
 143                 else
 144                         WRITE_ONCE(parent->rb_right, new);
 145         } else
 146                 WRITE_ONCE(root->rb_node, new);
 147 }
 148
 149 extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
 150         void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
 151
 152 static __always_inline struct rb_node *
 153 __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
 154                      struct rb_node **leftmost,
 155                      const struct rb_augment_callbacks *augment)
 156 {
 157         struct rb_node *child = node->rb_right;
 158         struct rb_node *tmp = node->rb_left;
 159         struct rb_node *parent, *rebalance;
 160         unsigned long pc;
 161
 162         if (leftmost && node == *leftmost)
 163                 *leftmost = rb_next(node);
 164
 165         if (!tmp) {
 166                 /*
 167                  * Case 1: node to erase has no more than 1 child (easy!)
 168                  *
 169                  * Note that if there is one child it must be red due to 5)
 170                  * and node must be black due to 4). We adjust colors locally
 171                  * so as to bypass __rb_erase_color() later on.
 172                  */
 173                 pc = node->__rb_parent_color;
 174                 parent = __rb_parent(pc);
 175                 __rb_change_child(node, child, parent, root);
 176                 if (child) {
 177                         child->__rb_parent_color = pc;
 178                         rebalance = NULL;
 179                 } else
 180                         rebalance = __rb_is_black(pc) ? parent : NULL;
 181                 tmp = parent;
 182         } else if (!child) {
 183                 /* Still case 1, but this time the child is node->rb_left */
 184                 tmp->__rb_parent_color = pc = node->__rb_parent_color;
 185                 parent = __rb_parent(pc);
 186                 __rb_change_child(node, tmp, parent, root);
 187                 rebalance = NULL;
 188                 tmp = parent;
 189         } else {
 190                 struct rb_node *successor = child, *child2;
 191
 192                 tmp = child->rb_left;
 193                 if (!tmp) {
 194                         /*
 195                          * Case 2: node's successor is its right child
 196                          *
 197                          *    (n)          (s)
 198                          *    / \          / \
 199                          *  (x) (s)  ->  (x) (c)
 200                          *        \
 201                          *        (c)
 202                          */
 203                         parent = successor;
 204                         child2 = successor->rb_right;
 205
 206                         augment->copy(node, successor);
 207                 } else {
 208                         /*
 209                          * Case 3: node's successor is leftmost under
 210                          * node's right child subtree
 211                          *
 212                          *    (n)          (s)
 213                          *    / \          / \
 214                          *  (x) (y)  ->  (x) (y)
 215                          *      /            /
 216                          *    (p)          (p)
 217                          *    /            /
 218                          *  (s)          (c)
 219                          *    \
 220                          *    (c)
 221                          */
 222                         do {
 223                                 parent = successor;
 224                                 successor = tmp;
 225                                 tmp = tmp->rb_left;
 226                         } while (tmp);
 227                         child2 = successor->rb_right;
 228                         WRITE_ONCE(parent->rb_left, child2);
 229                         WRITE_ONCE(successor->rb_right, child);
 230                         rb_set_parent(child, successor);
 231
 232                         augment->copy(node, successor);
 233                         augment->propagate(parent, successor);
 234                 }
 235
 236                 tmp = node->rb_left;
 237                 WRITE_ONCE(successor->rb_left, tmp);
 238                 rb_set_parent(tmp, successor);
 239
 240                 pc = node->__rb_parent_color;
 241                 tmp = __rb_parent(pc);
 242                 __rb_change_child(node, successor, tmp, root);
 243
 244                 if (child2) {
 245                         successor->__rb_parent_color = pc;
 246                         rb_set_parent_color(child2, parent, RB_BLACK);
 247                         rebalance = NULL;
 248                 } else {
 249                         unsigned long pc2 = successor->__rb_parent_color;
 250                         successor->__rb_parent_color = pc;
 251                         rebalance = __rb_is_black(pc2) ? parent : NULL;
 252                 }
 253                 tmp = successor;
 254         }
 255
 256         augment->propagate(tmp, NULL);
 257         return rebalance;
 258 }
 259
 260 static __always_inline void
 261 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
 262                    const struct rb_augment_callbacks *augment)
 263 {
 264         struct rb_node *rebalance = __rb_erase_augmented(node, root,
 265                                                          NULL, augment);
 266         if (rebalance)
 267                 __rb_erase_color(rebalance, root, augment->rotate);
 268 }
 269
 270 static __always_inline void
 271 rb_erase_augmented_cached(struct rb_node *node, struct rb_root_cached *root,
 272                           const struct rb_augment_callbacks *augment)
 273 {
 274         struct rb_node *rebalance = __rb_erase_augmented(node, &root->rb_root,
 275                                                          &root->rb_leftmost,
 276                                                          augment);
 277         if (rebalance)
 278                 __rb_erase_color(rebalance, &root->rb_root, augment->rotate);
 279 }
 280
 281 #endif /* _TOOLS_LINUX_RBTREE_AUGMENTED_H */