lib/rbtree.c

   1 /*
   2   Red Black Trees
   3   (C) 1999  Andrea Arcangeli <andrea@suse.de>
   4   (C) 2002  David Woodhouse <dwmw2@infradead.org>
   5
   6   This program is free software; you can redistribute it and/or modify
   7   it under the terms of the GNU General Public License as published by
   8   the Free Software Foundation; either version 2 of the License, or
   9   (at your option) any later version.
  10
  11   This program is distributed in the hope that it will be useful,
  12   but WITHOUT ANY WARRANTY; without even the implied warranty of
  13   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14   GNU General Public License for more details.
  15
  16   You should have received a copy of the GNU General Public License
  17   along with this program; if not, write to the Free Software
  18   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  19
  20   linux/lib/rbtree.c
  21 */
  22
  23 #include <linux/rbtree.h>
  24 #include <linux/export.h>
  25
  26 #define RB_RED          0
  27 #define RB_BLACK        1
  28
  29 #define rb_color(r)   ((r)->__rb_parent_color & 1)
  30 #define rb_is_red(r)   (!rb_color(r))
  31 #define rb_is_black(r) rb_color(r)
  32 #define rb_set_red(r)  do { (r)->__rb_parent_color &= ~1; } while (0)
  33 #define rb_set_black(r)  do { (r)->__rb_parent_color |= 1; } while (0)
  34
  35 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
  36 {
  37         rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
  38 }
  39 static inline void rb_set_color(struct rb_node *rb, int color)
  40 {
  41         rb->__rb_parent_color = (rb->__rb_parent_color & ~1) | color;
  42 }
  43
  44 static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
  45 {
  46         struct rb_node *right = node->rb_right;
  47         struct rb_node *parent = rb_parent(node);
  48
  49         if ((node->rb_right = right->rb_left))
  50                 rb_set_parent(right->rb_left, node);
  51         right->rb_left = node;
  52
  53         rb_set_parent(right, parent);
  54
  55         if (parent)
  56         {
  57                 if (node == parent->rb_left)
  58                         parent->rb_left = right;
  59                 else
  60                         parent->rb_right = right;
  61         }
  62         else
  63                 root->rb_node = right;
  64         rb_set_parent(node, right);
  65 }
  66
  67 static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
  68 {
  69         struct rb_node *left = node->rb_left;
  70         struct rb_node *parent = rb_parent(node);
  71
  72         if ((node->rb_left = left->rb_right))
  73                 rb_set_parent(left->rb_right, node);
  74         left->rb_right = node;
  75
  76         rb_set_parent(left, parent);
  77
  78         if (parent)
  79         {
  80                 if (node == parent->rb_right)
  81                         parent->rb_right = left;
  82                 else
  83                         parent->rb_left = left;
  84         }
  85         else
  86                 root->rb_node = left;
  87         rb_set_parent(node, left);
  88 }
  89
  90 void rb_insert_color(struct rb_node *node, struct rb_root *root)
  91 {
  92         struct rb_node *parent, *gparent;
  93
  94         while ((parent = rb_parent(node)) && rb_is_red(parent))
  95         {
  96                 gparent = rb_parent(parent);
  97
  98                 if (parent == gparent->rb_left)
  99                 {
 100                         {
 101                                 register struct rb_node *uncle = gparent->rb_right;
 102                                 if (uncle && rb_is_red(uncle))
 103                                 {
 104                                         rb_set_black(uncle);
 105                                         rb_set_black(parent);
 106                                         rb_set_red(gparent);
 107                                         node = gparent;
 108                                         continue;
 109                                 }
 110                         }
 111
 112                         if (parent->rb_right == node)
 113                         {
 114                                 register struct rb_node *tmp;
 115                                 __rb_rotate_left(parent, root);
 116                                 tmp = parent;
 117                                 parent = node;
 118                                 node = tmp;
 119                         }
 120
 121                         rb_set_black(parent);
 122                         rb_set_red(gparent);
 123                         __rb_rotate_right(gparent, root);
 124                 } else {
 125                         {
 126                                 register struct rb_node *uncle = gparent->rb_left;
 127                                 if (uncle && rb_is_red(uncle))
 128                                 {
 129                                         rb_set_black(uncle);
 130                                         rb_set_black(parent);
 131                                         rb_set_red(gparent);
 132                                         node = gparent;
 133                                         continue;
 134                                 }
 135                         }
 136
 137                         if (parent->rb_left == node)
 138                         {
 139                                 register struct rb_node *tmp;
 140                                 __rb_rotate_right(parent, root);
 141                                 tmp = parent;
 142                                 parent = node;
 143                                 node = tmp;
 144                         }
 145
 146                         rb_set_black(parent);
 147                         rb_set_red(gparent);
 148                         __rb_rotate_left(gparent, root);
 149                 }
 150         }
 151
 152         rb_set_black(root->rb_node);
 153 }
 154 EXPORT_SYMBOL(rb_insert_color);
 155
 156 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 157                              struct rb_root *root)
 158 {
 159         struct rb_node *other;
 160
 161         while ((!node || rb_is_black(node)) && node != root->rb_node)
 162         {
 163                 if (parent->rb_left == node)
 164                 {
 165                         other = parent->rb_right;
 166                         if (rb_is_red(other))
 167                         {
 168                                 rb_set_black(other);
 169                                 rb_set_red(parent);
 170                                 __rb_rotate_left(parent, root);
 171                                 other = parent->rb_right;
 172                         }
 173                         if ((!other->rb_left || rb_is_black(other->rb_left)) &&
 174                             (!other->rb_right || rb_is_black(other->rb_right)))
 175                         {
 176                                 rb_set_red(other);
 177                                 node = parent;
 178                                 parent = rb_parent(node);
 179                         }
 180                         else
 181                         {
 182                                 if (!other->rb_right || rb_is_black(other->rb_right))
 183                                 {
 184                                         rb_set_black(other->rb_left);
 185                                         rb_set_red(other);
 186                                         __rb_rotate_right(other, root);
 187                                         other = parent->rb_right;
 188                                 }
 189                                 rb_set_color(other, rb_color(parent));
 190                                 rb_set_black(parent);
 191                                 rb_set_black(other->rb_right);
 192                                 __rb_rotate_left(parent, root);
 193                                 node = root->rb_node;
 194                                 break;
 195                         }
 196                 }
 197                 else
 198                 {
 199                         other = parent->rb_left;
 200                         if (rb_is_red(other))
 201                         {
 202                                 rb_set_black(other);
 203                                 rb_set_red(parent);
 204                                 __rb_rotate_right(parent, root);
 205                                 other = parent->rb_left;
 206                         }
 207                         if ((!other->rb_left || rb_is_black(other->rb_left)) &&
 208                             (!other->rb_right || rb_is_black(other->rb_right)))
 209                         {
 210                                 rb_set_red(other);
 211                                 node = parent;
 212                                 parent = rb_parent(node);
 213                         }
 214                         else
 215                         {
 216                                 if (!other->rb_left || rb_is_black(other->rb_left))
 217                                 {
 218                                         rb_set_black(other->rb_right);
 219                                         rb_set_red(other);
 220                                         __rb_rotate_left(other, root);
 221                                         other = parent->rb_left;
 222                                 }
 223                                 rb_set_color(other, rb_color(parent));
 224                                 rb_set_black(parent);
 225                                 rb_set_black(other->rb_left);
 226                                 __rb_rotate_right(parent, root);
 227                                 node = root->rb_node;
 228                                 break;
 229                         }
 230                 }
 231         }
 232         if (node)
 233                 rb_set_black(node);
 234 }
 235
 236 void rb_erase(struct rb_node *node, struct rb_root *root)
 237 {
 238         struct rb_node *child, *parent;
 239         int color;
 240
 241         if (!node->rb_left)
 242                 child = node->rb_right;
 243         else if (!node->rb_right)
 244                 child = node->rb_left;
 245         else
 246         {
 247                 struct rb_node *old = node, *left;
 248
 249                 node = node->rb_right;
 250                 while ((left = node->rb_left) != NULL)
 251                         node = left;
 252
 253                 if (rb_parent(old)) {
 254                         if (rb_parent(old)->rb_left == old)
 255                                 rb_parent(old)->rb_left = node;
 256                         else
 257                                 rb_parent(old)->rb_right = node;
 258                 } else
 259                         root->rb_node = node;
 260
 261                 child = node->rb_right;
 262                 parent = rb_parent(node);
 263                 color = rb_color(node);
 264
 265                 if (parent == old) {
 266                         parent = node;
 267                 } else {
 268                         if (child)
 269                                 rb_set_parent(child, parent);
 270                         parent->rb_left = child;
 271
 272                         node->rb_right = old->rb_right;
 273                         rb_set_parent(old->rb_right, node);
 274                 }
 275
 276                 node->__rb_parent_color = old->__rb_parent_color;
 277                 node->rb_left = old->rb_left;
 278                 rb_set_parent(old->rb_left, node);
 279
 280                 goto color;
 281         }
 282
 283         parent = rb_parent(node);
 284         color = rb_color(node);
 285
 286         if (child)
 287                 rb_set_parent(child, parent);
 288         if (parent)
 289         {
 290                 if (parent->rb_left == node)
 291                         parent->rb_left = child;
 292                 else
 293                         parent->rb_right = child;
 294         }
 295         else
 296                 root->rb_node = child;
 297
 298  color:
 299         if (color == RB_BLACK)
 300                 __rb_erase_color(child, parent, root);
 301 }
 302 EXPORT_SYMBOL(rb_erase);
 303
 304 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
 305 {
 306         struct rb_node *parent;
 307
 308 up:
 309         func(node, data);
 310         parent = rb_parent(node);
 311         if (!parent)
 312                 return;
 313
 314         if (node == parent->rb_left && parent->rb_right)
 315                 func(parent->rb_right, data);
 316         else if (parent->rb_left)
 317                 func(parent->rb_left, data);
 318
 319         node = parent;
 320         goto up;
 321 }
 322
 323 /*
 324  * after inserting @node into the tree, update the tree to account for
 325  * both the new entry and any damage done by rebalance
 326  */
 327 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
 328 {
 329         if (node->rb_left)
 330                 node = node->rb_left;
 331         else if (node->rb_right)
 332                 node = node->rb_right;
 333
 334         rb_augment_path(node, func, data);
 335 }
 336 EXPORT_SYMBOL(rb_augment_insert);
 337
 338 /*
 339  * before removing the node, find the deepest node on the rebalance path
 340  * that will still be there after @node gets removed
 341  */
 342 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
 343 {
 344         struct rb_node *deepest;
 345
 346         if (!node->rb_right && !node->rb_left)
 347                 deepest = rb_parent(node);
 348         else if (!node->rb_right)
 349                 deepest = node->rb_left;
 350         else if (!node->rb_left)
 351                 deepest = node->rb_right;
 352         else {
 353                 deepest = rb_next(node);
 354                 if (deepest->rb_right)
 355                         deepest = deepest->rb_right;
 356                 else if (rb_parent(deepest) != node)
 357                         deepest = rb_parent(deepest);
 358         }
 359
 360         return deepest;
 361 }
 362 EXPORT_SYMBOL(rb_augment_erase_begin);
 363
 364 /*
 365  * after removal, update the tree to account for the removed entry
 366  * and any rebalance damage.
 367  */
 368 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
 369 {
 370         if (node)
 371                 rb_augment_path(node, func, data);
 372 }
 373 EXPORT_SYMBOL(rb_augment_erase_end);
 374
 375 /*
 376  * This function returns the first node (in sort order) of the tree.
 377  */
 378 struct rb_node *rb_first(const struct rb_root *root)
 379 {
 380         struct rb_node  *n;
 381
 382         n = root->rb_node;
 383         if (!n)
 384                 return NULL;
 385         while (n->rb_left)
 386                 n = n->rb_left;
 387         return n;
 388 }
 389 EXPORT_SYMBOL(rb_first);
 390
 391 struct rb_node *rb_last(const struct rb_root *root)
 392 {
 393         struct rb_node  *n;
 394
 395         n = root->rb_node;
 396         if (!n)
 397                 return NULL;
 398         while (n->rb_right)
 399                 n = n->rb_right;
 400         return n;
 401 }
 402 EXPORT_SYMBOL(rb_last);
 403
 404 struct rb_node *rb_next(const struct rb_node *node)
 405 {
 406         struct rb_node *parent;
 407
 408         if (RB_EMPTY_NODE(node))
 409                 return NULL;
 410
 411         /* If we have a right-hand child, go down and then left as far
 412            as we can. */
 413         if (node->rb_right) {
 414                 node = node->rb_right;
 415                 while (node->rb_left)
 416                         node=node->rb_left;
 417                 return (struct rb_node *)node;
 418         }
 419
 420         /* No right-hand children.  Everything down and left is
 421            smaller than us, so any 'next' node must be in the general
 422            direction of our parent. Go up the tree; any time the
 423            ancestor is a right-hand child of its parent, keep going
 424            up. First time it's a left-hand child of its parent, said
 425            parent is our 'next' node. */
 426         while ((parent = rb_parent(node)) && node == parent->rb_right)
 427                 node = parent;
 428
 429         return parent;
 430 }
 431 EXPORT_SYMBOL(rb_next);
 432
 433 struct rb_node *rb_prev(const struct rb_node *node)
 434 {
 435         struct rb_node *parent;
 436
 437         if (RB_EMPTY_NODE(node))
 438                 return NULL;
 439
 440         /* If we have a left-hand child, go down and then right as far
 441            as we can. */
 442         if (node->rb_left) {
 443                 node = node->rb_left;
 444                 while (node->rb_right)
 445                         node=node->rb_right;
 446                 return (struct rb_node *)node;
 447         }
 448
 449         /* No left-hand children. Go up till we find an ancestor which
 450            is a right-hand child of its parent */
 451         while ((parent = rb_parent(node)) && node == parent->rb_left)
 452                 node = parent;
 453
 454         return parent;
 455 }
 456 EXPORT_SYMBOL(rb_prev);
 457
 458 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 459                      struct rb_root *root)
 460 {
 461         struct rb_node *parent = rb_parent(victim);
 462
 463         /* Set the surrounding nodes to point to the replacement */
 464         if (parent) {
 465                 if (victim == parent->rb_left)
 466                         parent->rb_left = new;
 467                 else
 468                         parent->rb_right = new;
 469         } else {
 470                 root->rb_node = new;
 471         }
 472         if (victim->rb_left)
 473                 rb_set_parent(victim->rb_left, new);
 474         if (victim->rb_right)
 475                 rb_set_parent(victim->rb_right, new);
 476
 477         /* Copy the pointers/colour from the victim to the replacement */
 478         *new = *victim;
 479 }
 480 EXPORT_SYMBOL(rb_replace_node);