rbtree.c

   1 /*
   2   Red Black Trees
   3   (C) 1999  Andrea Arcangeli <andrea@suse.de>
   4   (C) 2002  David Woodhouse <dwmw2@infradead.org>
   5   (C) 2007  Andre Noll <maan@systemlinux.org>
   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   linux/lib/rbtree.c
  22 */
  23
  24 /** \file rbtree.c Red-black tree implementation. */
  25
  26 #include "stddef.h"
  27 #include "rbtree.h"
  28
  29 static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
  30 {
  31         struct rb_node *right = node->rb_right;
  32         struct rb_node *parent = rb_parent(node);
  33
  34         if ((node->rb_right = right->rb_left))
  35                 rb_set_parent(right->rb_left, node);
  36         right->rb_left = node;
  37
  38         rb_set_parent(right, parent);
  39
  40         if (parent)
  41         {
  42                 if (node == parent->rb_left)
  43                         parent->rb_left = right;
  44                 else
  45                         parent->rb_right = right;
  46         }
  47         else
  48                 root->rb_node = right;
  49         rb_set_parent(node, right);
  50         right->size = node->size;
  51         node->size = 1;
  52         if (node->rb_right)
  53                 node->size += node->rb_right->size;
  54         if (node->rb_left)
  55                 node->size += node->rb_left->size;
  56 }
  57
  58 static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
  59 {
  60         struct rb_node *left = node->rb_left;
  61         struct rb_node *parent = rb_parent(node);
  62
  63         if ((node->rb_left = left->rb_right))
  64                 rb_set_parent(left->rb_right, node);
  65         left->rb_right = node;
  66
  67         rb_set_parent(left, parent);
  68
  69         if (parent)
  70         {
  71                 if (node == parent->rb_right)
  72                         parent->rb_right = left;
  73                 else
  74                         parent->rb_left = left;
  75         }
  76         else
  77                 root->rb_node = left;
  78         rb_set_parent(node, left);
  79         left->size = node->size;
  80         node->size = 1;
  81         if (node->rb_right)
  82                 node->size += node->rb_right->size;
  83         if (node->rb_left)
  84                 node->size += node->rb_left->size;
  85 }
  86
  87 void rb_insert_color(struct rb_node *node, struct rb_root *root)
  88 {
  89         struct rb_node *parent, *gparent;
  90
  91         while ((parent = rb_parent(node)) && rb_is_red(parent))
  92         {
  93                 gparent = rb_parent(parent);
  94
  95                 if (parent == gparent->rb_left)
  96                 {
  97                         {
  98                                 register struct rb_node *uncle = gparent->rb_right;
  99                                 if (uncle && rb_is_red(uncle))
 100                                 {
 101                                         rb_set_black(uncle);
 102                                         rb_set_black(parent);
 103                                         rb_set_red(gparent);
 104                                         node = gparent;
 105                                         continue;
 106                                 }
 107                         }
 108
 109                         if (parent->rb_right == node)
 110                         {
 111                                 register struct rb_node *tmp;
 112                                 __rb_rotate_left(parent, root);
 113                                 tmp = parent;
 114                                 parent = node;
 115                                 node = tmp;
 116                         }
 117
 118                         rb_set_black(parent);
 119                         rb_set_red(gparent);
 120                         __rb_rotate_right(gparent, root);
 121                 } else {
 122                         {
 123                                 register struct rb_node *uncle = gparent->rb_left;
 124                                 if (uncle && rb_is_red(uncle))
 125                                 {
 126                                         rb_set_black(uncle);
 127                                         rb_set_black(parent);
 128                                         rb_set_red(gparent);
 129                                         node = gparent;
 130                                         continue;
 131                                 }
 132                         }
 133
 134                         if (parent->rb_left == node)
 135                         {
 136                                 register struct rb_node *tmp;
 137                                 __rb_rotate_right(parent, root);
 138                                 tmp = parent;
 139                                 parent = node;
 140                                 node = tmp;
 141                         }
 142
 143                         rb_set_black(parent);
 144                         rb_set_red(gparent);
 145                         __rb_rotate_left(gparent, root);
 146                 }
 147         }
 148
 149         rb_set_black(root->rb_node);
 150 }
 151
 152 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 153                              struct rb_root *root)
 154 {
 155         struct rb_node *other;
 156
 157         while ((!node || rb_is_black(node)) && node != root->rb_node)
 158         {
 159                 if (parent->rb_left == node)
 160                 {
 161                         other = parent->rb_right;
 162                         if (rb_is_red(other))
 163                         {
 164                                 rb_set_black(other);
 165                                 rb_set_red(parent);
 166                                 __rb_rotate_left(parent, root);
 167                                 other = parent->rb_right;
 168                         }
 169                         if ((!other->rb_left || rb_is_black(other->rb_left)) &&
 170                             (!other->rb_right || rb_is_black(other->rb_right)))
 171                         {
 172                                 rb_set_red(other);
 173                                 node = parent;
 174                                 parent = rb_parent(node);
 175                         }
 176                         else
 177                         {
 178                                 if (!other->rb_right || rb_is_black(other->rb_right))
 179                                 {
 180                                         struct rb_node *o_left;
 181                                         if ((o_left = other->rb_left))
 182                                                 rb_set_black(o_left);
 183                                         rb_set_red(other);
 184                                         __rb_rotate_right(other, root);
 185                                         other = parent->rb_right;
 186                                 }
 187                                 rb_set_color(other, rb_color(parent));
 188                                 rb_set_black(parent);
 189                                 if (other->rb_right)
 190                                         rb_set_black(other->rb_right);
 191                                 __rb_rotate_left(parent, root);
 192                                 node = root->rb_node;
 193                                 break;
 194                         }
 195                 }
 196                 else
 197                 {
 198                         other = parent->rb_left;
 199                         if (rb_is_red(other))
 200                         {
 201                                 rb_set_black(other);
 202                                 rb_set_red(parent);
 203                                 __rb_rotate_right(parent, root);
 204                                 other = parent->rb_left;
 205                         }
 206                         if ((!other->rb_left || rb_is_black(other->rb_left)) &&
 207                             (!other->rb_right || rb_is_black(other->rb_right)))
 208                         {
 209                                 rb_set_red(other);
 210                                 node = parent;
 211                                 parent = rb_parent(node);
 212                         }
 213                         else
 214                         {
 215                                 if (!other->rb_left || rb_is_black(other->rb_left))
 216                                 {
 217                                         register struct rb_node *o_right;
 218                                         if ((o_right = other->rb_right))
 219                                                 rb_set_black(o_right);
 220                                         rb_set_red(other);
 221                                         __rb_rotate_left(other, root);
 222                                         other = parent->rb_left;
 223                                 }
 224                                 rb_set_color(other, rb_color(parent));
 225                                 rb_set_black(parent);
 226                                 if (other->rb_left)
 227                                         rb_set_black(other->rb_left);
 228                                 __rb_rotate_right(parent, root);
 229                                 node = root->rb_node;
 230                                 break;
 231                         }
 232                 }
 233         }
 234         if (node)
 235                 rb_set_black(node);
 236 }
 237
 238 void rb_erase(struct rb_node *node, struct rb_root *root)
 239 {
 240         struct rb_node *child, *parent;
 241         int color;
 242
 243         if (!node->rb_left)
 244                 child = node->rb_right;
 245         else if (!node->rb_right)
 246                 child = node->rb_left;
 247         else
 248         {
 249                 struct rb_node *old = node, *left;
 250
 251                 node = node->rb_right;
 252                 while ((left = node->rb_left) != NULL)
 253                         node = left;
 254                 child = node->rb_right;
 255                 parent = rb_parent(node);
 256                 color = rb_color(node);
 257
 258                 if (child)
 259                         rb_set_parent(child, parent);
 260                 if (parent == old) {
 261                         parent->rb_right = child;
 262                         parent = node;
 263                 } else
 264                         parent->rb_left = child;
 265
 266                 node->rb_parent_color = old->rb_parent_color;
 267                 node->rb_right = old->rb_right;
 268                 node->rb_left = old->rb_left;
 269                 node->size = old->size;
 270
 271                 if (rb_parent(old))
 272                 {
 273                         if (rb_parent(old)->rb_left == old)
 274                                 rb_parent(old)->rb_left = node;
 275                         else
 276                                 rb_parent(old)->rb_right = node;
 277                 } else
 278                         root->rb_node = node;
 279
 280                 rb_set_parent(old->rb_left, node);
 281                 if (old->rb_right)
 282                         rb_set_parent(old->rb_right, node);
 283                 goto color;
 284         }
 285
 286         parent = rb_parent(node);
 287         color = rb_color(node);
 288
 289         if (child)
 290                 rb_set_parent(child, parent);
 291         if (parent)
 292         {
 293                 if (parent->rb_left == node)
 294                         parent->rb_left = child;
 295                 else
 296                         parent->rb_right = child;
 297         }
 298         else
 299                 root->rb_node = child;
 300
 301  color:
 302         if (color == RB_BLACK)
 303                 __rb_erase_color(child, parent, root);
 304 }
 305
 306 /*
 307  * This function returns the first node (in sort order) of the tree.
 308  */
 309 struct rb_node *rb_first(struct rb_root *root)
 310 {
 311         struct rb_node  *n;
 312
 313         n = root->rb_node;
 314         if (!n)
 315                 return NULL;
 316         while (n->rb_left)
 317                 n = n->rb_left;
 318         return n;
 319 }
 320
 321 struct rb_node *rb_last(struct rb_root *root)
 322 {
 323         struct rb_node  *n;
 324
 325         n = root->rb_node;
 326         if (!n)
 327                 return NULL;
 328         while (n->rb_right)
 329                 n = n->rb_right;
 330         return n;
 331 }
 332
 333 struct rb_node *rb_next(struct rb_node *node)
 334 {
 335         struct rb_node *parent;
 336
 337         if (rb_parent(node) == node)
 338                 return NULL;
 339
 340         /* If we have a right-hand child, go down and then left as far
 341            as we can. */
 342         if (node->rb_right) {
 343                 node = node->rb_right;
 344                 while (node->rb_left)
 345                         node=node->rb_left;
 346                 return node;
 347         }
 348
 349         /* No right-hand children.  Everything down and left is
 350            smaller than us, so any 'next' node must be in the general
 351            direction of our parent. Go up the tree; any time the
 352            ancestor is a right-hand child of its parent, keep going
 353            up. First time it's a left-hand child of its parent, said
 354            parent is our 'next' node. */
 355         while ((parent = rb_parent(node)) && node == parent->rb_right)
 356                 node = parent;
 357
 358         return parent;
 359 }
 360
 361 struct rb_node *rb_prev(struct rb_node *node)
 362 {
 363         struct rb_node *parent;
 364
 365         if (rb_parent(node) == node)
 366                 return NULL;
 367
 368         /* If we have a left-hand child, go down and then right as far
 369            as we can. */
 370         if (node->rb_left) {
 371                 node = node->rb_left;
 372                 while (node->rb_right)
 373                         node=node->rb_right;
 374                 return node;
 375         }
 376
 377         /* No left-hand children. Go up till we find an ancestor which
 378            is a right-hand child of its parent */
 379         while ((parent = rb_parent(node)) && node == parent->rb_left)
 380                 node = parent;
 381
 382         return parent;
 383 }
 384
 385 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 386                      struct rb_root *root)
 387 {
 388         struct rb_node *parent = rb_parent(victim);
 389
 390         /* Set the surrounding nodes to point to the replacement */
 391         if (parent) {
 392                 if (victim == parent->rb_left)
 393                         parent->rb_left = new;
 394                 else
 395                         parent->rb_right = new;
 396         } else {
 397                 root->rb_node = new;
 398         }
 399         if (victim->rb_left)
 400                 rb_set_parent(victim->rb_left, new);
 401         if (victim->rb_right)
 402                 rb_set_parent(victim->rb_right, new);
 403
 404         /* Copy the pointers/colour from the victim to the replacement */
 405         *new = *victim;
 406 }
 407
 408 /**
 409  * Get the n-th node (in sort order) of the tree.
 410  *
 411  * \param node The root of the subtree to consider.
 412  * \param n The order statistic to compute.
 413  *
 414  * \return Pointer to the \a n th greatest node on success, \p NULL on errors.
 415  */
 416 struct rb_node *rb_nth(struct rb_node *node, unsigned n)
 417 {
 418         unsigned size = 1;
 419
 420         if (!node)
 421                 return NULL;
 422         if (node->rb_left)
 423                 size += node->rb_left->size;
 424         if (n == size)
 425                 return node;
 426         if (n < size)
 427                 return rb_nth(node->rb_left, n);
 428         return rb_nth(node->rb_right, n - size);
 429 }
 430
 431 /**
 432  * Get the rank of a node in O(log n) time.
 433  *
 434  * \param node The node to get the rank of.
 435  * \param rank Result pointer.
 436  *
 437  * \return Positive on success, -1 on errors.
 438  */
 439 int rb_rank(struct rb_node *node, unsigned *rank)
 440 {
 441         *rank = 1;
 442         struct rb_node *parent;
 443
 444         if (!node)
 445                 return -1;
 446         if (node->rb_left)
 447                 *rank += node->rb_left->size;
 448         while ((parent = rb_parent(node))) {
 449                 if (node == parent->rb_right) {
 450                         (*rank)++;
 451                         if (parent->rb_left)
 452                                 *rank += parent->rb_left->size;
 453                 }
 454                 node = parent;
 455         }
 456         return 1;
 457 }