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