include/c++/util/rbtree.h

   1 // util/rbtree.h -- A red/black tree implementation
   2 //
   3 // This software is copyright (c) 2006 Scott Wood <scott@buserror.net>.
   4 //
   5 // Permission is hereby granted, free of charge, to any person obtaining a copy of
   6 // this software and associated documentation files (the "Software"), to deal with
   7 // the Software without restriction, including without limitation the rights to
   8 // use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
   9 // of the Software, and to permit persons to whom the Software is furnished to do
  10 // so, subject to the following conditions:
  11 //
  12 //     * Redistributions of source code must retain the above copyright notice,
  13 //       this list of conditions and the following disclaimers.
  14 //
  15 //     * Redistributions in binary form must reproduce the above copyright notice,
  16 //       this list of conditions and the following disclaimers in the
  17 //       documentation and/or other materials provided with the distribution.
  18 //
  19 //     * The names of the Software's authors and/or contributors
  20 //       may not be used to endorse or promote products derived from
  21 //       this Software without specific prior written permission.
  22 //
  23 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  24 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
  25 // FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
  26 // CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  27 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  28 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE
  29 // SOFTWARE.
  30
  31 #ifndef _UTIL_RBTREE_H
  32 #define _UTIL_RBTREE_H
  33
  34 #include <assert.h>
  35 #include <stddef.h>
  36 #include <stdint.h>
  37
  38 namespace Util {
  39         // T must have an RBTree<T, NodeV, KeyV>::Node member called rbtree_node.
  40         // NodeV < NodeV, NodeV < KeyV, and NodeV > KeyV must be supported.
  41         //
  42         // Using NodeV != KeyV allows things such as having NodeV represent a
  43         // range of values and KeyV a single value, causing find() to search
  44         // for the node in whose range KeyV falls.  It is an error to add
  45         // nodes that are not well ordered.
  46
  47         template <typename T, typename NodeV, typename KeyV>
  48         class RBTree {
  49         public:
  50                 struct Node {
  51                         Node *parent, *left, *right;
  52                         NodeV value;
  53                         bool red;
  54
  55                         // A pointer to the parent's left or right pointer
  56                         // (whichever side this node is on).  This is also
  57                         // used for sanity checks (it is NULL for nodes not
  58                         // on a tree).
  59
  60                         Node **parentlink;
  61
  62                         Node()
  63                         {
  64                                 parent = left = right = NULL;
  65                                 parentlink = NULL;
  66                         }
  67
  68                         bool is_on_rbtree()
  69                         {
  70                                 return parentlink != NULL;
  71                         }
  72                 };
  73
  74         private:
  75                 Node *top;
  76
  77                 T *node_to_type(Node *n)
  78                 {
  79                         return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(n) -
  80                                                      offsetof(T, rbtree_node));
  81                 }
  82
  83                 void assert_rb_condition(Node *left, Node *right)
  84                 {
  85                         assert(left->value < right->value);
  86                         assert(!(right->value < left->value));
  87                         assert(!left->red || !right->red);
  88                 }
  89
  90                 void assert_well_ordered(Node *left, NodeV right)
  91                 {
  92                         assert(left->value < right);
  93                         assert(!(left->value > right));
  94                         assert(!left->red || !right->red);
  95                 }
  96
  97                 Node *find_node(Node *n, NodeV &val, Node ***link);
  98                 Node *find_node_key(KeyV val, bool exact = true);
  99
 100                 bool node_is_left(Node *n)
 101                 {
 102                         assert(n->parentlink == &n->parent->left ||
 103                                n->parentlink == &n->parent->right);
 104
 105                         return n->parentlink == &n->parent->left;
 106                 }
 107
 108                 Node *sibling(Node *n)
 109                 {
 110                         if (node_is_left(n))
 111                                 return n->parent->right;
 112
 113                         return n->parent->left;
 114                 }
 115
 116                 Node *uncle(Node *n)
 117                 {
 118                         return sibling(n->parent);
 119                 }
 120
 121                 void rotate_left(Node *n);
 122                 void rotate_right(Node *n);
 123
 124                 // B must have no right child and have A as an ancestor.
 125                 void swap(Node *a, Node *b);
 126
 127                 bool red(Node *n)
 128                 {
 129                         return n && n->red;
 130                 }
 131
 132 public:
 133                 RBTree()
 134                 {
 135                         top = NULL;
 136                 }
 137
 138                 bool empty()
 139                 {
 140                         return top == NULL;
 141                 }
 142
 143                 T *find(KeyV value)
 144                 {
 145                         Node *n = find_node_key(value);
 146
 147                         if (n)
 148                                 return node_to_type(n);
 149
 150                         return NULL;
 151                 }
 152
 153                 T *find_nearest(KeyV value)
 154                 {
 155                         Node *n = find_node_key(value, false);
 156
 157                         if (n)
 158                                 return node_to_type(n);
 159
 160                         // Should only happen on an empty tree
 161                         return NULL;
 162                 }
 163
 164                 void add(T *t);
 165                 void del(T *t);
 166         };
 167
 168         template <typename T, typename NodeV, typename KeyV>
 169         typename RBTree<T, NodeV, KeyV>::Node *
 170         RBTree<T, NodeV, KeyV>::find_node(Node *n, NodeV &val, Node ***link)
 171         {
 172                 if (link) {
 173                         if (n == top)
 174                                 *link = &top;
 175                         else
 176                                 *link = n->parentlink;
 177                 }
 178
 179                 if (!top)
 180                         return NULL;
 181
 182                 while (true) {
 183                         assert(n);
 184                         Node *next;
 185
 186                         if (n->value < val) {
 187                                 next = n->right;
 188
 189                                 if (next) {
 190                                         assert_rb_condition(n, next);
 191                                 } else {
 192                                         if (link)
 193                                                 *link = &n->right;
 194
 195                                         break;
 196                                 }
 197                         } else {
 198                                 // This assert detects duplicate nodes, but not
 199                                 // overlapping ranges (or otherwise non-well-ordered
 200                                 // values).
 201
 202                                 assert(val < n->value);
 203                                 next = n->left;
 204
 205                                 if (next) {
 206                                         assert_rb_condition(next, n);
 207                                 } else {
 208                                         if (link)
 209                                                 *link = &n->left;
 210
 211                                         break;
 212                                 }
 213                         }
 214
 215                         n = next;
 216                 }
 217
 218                 return n;
 219         }
 220
 221         template <typename T, typename NodeV, typename KeyV>
 222         typename RBTree<T, NodeV, KeyV>::Node *
 223         RBTree<T, NodeV, KeyV>::find_node_key(KeyV val, bool exact)
 224         {
 225                 Node *n = top;
 226
 227                 while (n) {
 228                         Node *next;
 229
 230                         if (n->value < val) {
 231                                 next = n->right;
 232
 233                                 if (next)
 234                                         assert_rb_condition(n, next);
 235                         } else if (n->value > val) {
 236                                 next = n->left;
 237
 238                                 if (next)
 239                                         assert_rb_condition(next, n);
 240                         } else {
 241                                 break;
 242                         }
 243
 244                         if (!next && !exact)
 245                                 return n;
 246
 247                         n = next;
 248                 }
 249
 250                 return n;
 251         }
 252
 253         template <typename T, typename NodeV, typename KeyV>
 254         void RBTree<T, NodeV, KeyV>::rotate_left(Node *n)
 255         {
 256                 Node *new_top = n->right;
 257
 258                 assert(*n->parentlink == n);
 259                 *n->parentlink = new_top;
 260
 261                 assert(new_top->parent == n);
 262                 assert(new_top->parentlink == &n->right);
 263
 264                 new_top->parent = n->parent;
 265                 new_top->parentlink = n->parentlink;
 266
 267                 n->parent = new_top;
 268                 n->parentlink = &new_top->left;
 269
 270                 n->right = new_top->left;
 271                 new_top->left = n;
 272
 273                 if (n->right) {
 274                         assert(n->right->parent == new_top);
 275                         assert(n->right->parentlink == &new_top->left);
 276
 277                         n->right->parent = n;
 278                         n->right->parentlink = &n->right;
 279                 }
 280         }
 281
 282         template <typename T, typename NodeV, typename KeyV>
 283         void RBTree<T, NodeV, KeyV>::rotate_right(Node *n)
 284         {
 285                 Node *new_top = n->left;
 286
 287                 assert(*n->parentlink == n);
 288                 *n->parentlink = new_top;
 289
 290                 assert(new_top->parent == n);
 291                 assert(new_top->parentlink == &n->left);
 292
 293                 new_top->parent = n->parent;
 294                 new_top->parentlink = n->parentlink;
 295
 296                 n->parent = new_top;
 297                 n->parentlink = &new_top->right;
 298
 299                 n->left = new_top->right;
 300                 new_top->right = n;
 301
 302                 if (n->left) {
 303                         assert(n->left->parent == new_top);
 304                         assert(n->left->parentlink == &new_top->right);
 305
 306                         n->left->parent = n;
 307                         n->left->parentlink = &n->left;
 308                 }
 309         }
 310
 311         // B must have no right child and have A as an ancestor.
 312         template <typename T, typename NodeV, typename KeyV>
 313         void RBTree<T, NodeV, KeyV>::swap(Node *a, Node *b)
 314         {
 315                 Node *bparent = b->parent;
 316                 Node **bparentlink = b->parentlink;
 317
 318                 assert(!b->right);
 319                 assert(a->left || a->right);
 320
 321                 b->parent = a->parent;
 322                 b->parentlink = a->parentlink;
 323
 324                 if (bparent == a) {
 325                         a->parent = b;
 326                         a->parentlink = &b->left;
 327                 } else {
 328                         a->parent = bparent;
 329                         a->parentlink = bparentlink;
 330                 }
 331
 332                 assert(a->parent != a);
 333                 assert(b->parent != b);
 334
 335                 Node *bleft = b->left;
 336                 b->left = a->left;
 337                 a->left = bleft;
 338
 339                 b->right = a->right;
 340                 a->right = NULL;
 341
 342                 *a->parentlink = a;
 343                 *b->parentlink = b;
 344
 345                 assert(a->parent != a);
 346                 assert(b->parent != b);
 347
 348                 bool bred = b->red;
 349                 b->red = a->red;
 350                 a->red = bred;
 351
 352                 if (a->left) {
 353                         a->left->parent = a;
 354                         a->left->parentlink = &a->left;
 355                 }
 356
 357                 if (b->right) {
 358                         b->right->parent = b;
 359                         b->right->parentlink = &b->right;
 360                 }
 361
 362                 assert(b->left);
 363                 b->left->parent = b;
 364                 b->left->parentlink = &b->left;
 365         }
 366
 367         template <typename T, typename NodeV, typename KeyV>
 368         void RBTree<T, NodeV, KeyV>::add(T *t)
 369         {
 370                 Node *n = &t->rbtree_node;
 371                 assert(!n->is_on_rbtree());
 372
 373                 Node *insert_at = find_node(top, n->value, &n->parentlink);
 374
 375                 assert(insert_at || !top);
 376                 *n->parentlink = n;
 377                 n->parent = insert_at;
 378                 n->left = n->right = NULL;
 379
 380         repeat:
 381                 assert(n->parentlink);
 382                 assert(*n->parentlink == n);
 383
 384                 if (!n->parent) {
 385                         n->red = false;
 386                         return;
 387                 }
 388
 389                 assert(n->parent->value < n->value != n->value < n->parent->value);
 390                 n->red = true;
 391
 392                 if (!n->parent->red)
 393                         return;
 394
 395                 Node *unc = uncle(n);
 396                 if (red(unc)) {
 397                         assert(!n->parent->parent->red);
 398                         n->parent->red = unc->red = false;
 399                         n = n->parent->parent;
 400                         goto repeat;
 401                 }
 402
 403                 if (node_is_left(n)) {
 404                         if (!node_is_left(n->parent)) {
 405                                 rotate_right(n->parent);
 406                                 n = n->right;
 407                         }
 408                 } else {
 409                         if (node_is_left(n->parent)) {
 410                                 rotate_left(n->parent);
 411                                 n = n->left;
 412                         }
 413                 }
 414
 415                 assert(n->parent->red);
 416                 assert(!red(uncle(n)));
 417                 assert(!n->parent->parent->red);
 418
 419                 n->parent->red = false;
 420                 n->parent->parent->red = true;
 421
 422                 if (node_is_left(n)) {
 423                         assert(node_is_left(n->parent));
 424                         rotate_right(n->parent->parent);
 425                 } else {
 426                         assert(!node_is_left(n->parent));
 427                         rotate_left(n->parent->parent);
 428                 }
 429         }
 430
 431         template <typename T, typename NodeV, typename KeyV>
 432         void RBTree<T, NodeV, KeyV>::del(T *t)
 433         {
 434                 Node *n = &t->rbtree_node;
 435                 assert(*n->parentlink == n);
 436
 437                 if (n->left && n->right) {
 438                         Node *highest_on_left = find_node(n->left, n->value, NULL);
 439                         assert(!highest_on_left->right);
 440                         swap(n, highest_on_left);
 441                         assert(!n->right);
 442                 }
 443
 444                 Node *parent = n->parent;
 445                 Node *child = n->left ? n->left : n->right;
 446
 447                 if (child) {
 448                         assert(child->parent == n);
 449
 450                         child->parent = n->parent;
 451                         child->parentlink = n->parentlink;
 452                         *child->parentlink = child;
 453                         assert(child != parent);
 454                 } else {
 455                         *n->parentlink = NULL;
 456                 }
 457
 458                 n->parentlink = NULL;
 459
 460                 if (n->red)
 461                         return;
 462
 463                 n = child;
 464
 465                 if (red(n)) {
 466                         n->red = false;
 467                         return;
 468                 }
 469
 470         repeat:
 471                 if (n == top) {
 472                         assert(!red(n));
 473                         return;
 474                 }
 475
 476                 Node *sib;
 477
 478                 if (n)
 479                         sib = sibling(n);
 480                 else
 481                         sib = parent->left ? parent->left : parent->right;
 482
 483                 if (sib->red) {
 484                         assert(!parent->red);
 485                         assert(!red(sib->left));
 486                         assert(!red(sib->right));
 487
 488                         parent->red = true;
 489                         sib->red = false;
 490
 491                         if (node_is_left(sib)) {
 492                                 rotate_right(parent);
 493                                 sib = parent->left;
 494                         } else {
 495                                 rotate_left(parent);
 496                                 sib = parent->right;
 497                         }
 498
 499                         if (n)
 500                                 assert(sib == sibling(n));
 501                 } else if (!parent->red && !red(sib->left) && !red(sib->right)) {
 502                         sib->red = true;
 503                         assert(n != parent);
 504                         n = parent;
 505                         parent = parent->parent;
 506                         goto repeat;
 507                 }
 508
 509                 assert(!sib->red);
 510
 511                 if (!parent->red && !red(sib->left) && !red(sib->right)) {
 512                         sib->red = true;
 513                         parent->red = false;
 514                         return;
 515                 }
 516
 517                 if (!red(sib->left) && !red(sib->right)) {
 518                         assert(parent->red);
 519                         sib->red = true;
 520                         parent->red = false;
 521                         return;
 522                 }
 523
 524                 if (node_is_left(sib) && !red(sib->left) && red(sib->right)) {
 525                         sib->red = true;
 526                         sib->right->red = false;
 527                         rotate_left(sib);
 528                         sib = sib->parent;
 529                 } else if (!node_is_left(sib) && !red(sib->right) && red(sib->left)) {
 530                         sib->red = true;
 531                         sib->left->red = false;
 532                         rotate_right(sib);
 533                         sib = sib->parent;
 534                 }
 535
 536                 assert(parent == sib->parent);
 537                 assert(!sib->red);
 538
 539                 sib->red = parent->red;
 540                 parent->red = false;
 541
 542                 if (node_is_left(sib)) {
 543                         assert(sib->left->red);
 544                         sib->left->red = false;
 545                         rotate_right(parent);
 546                 } else {
 547                         assert(sib->right->red);
 548                         sib->right->red = false;
 549                         rotate_left(parent);
 550                 }
 551         }
 552 }
 553
 554 #endif