Source/WTF/wtf/RedBlackTree.h

   1 /*
   2  * Copyright (C) 2010, 2011 Apple Inc. All rights reserved.
   3  *
   4  * Redistribution and use in source and binary forms, with or without
   5  * modification, are permitted provided that the following conditions
   6  * are met:
   7  *
   8  * 1.  Redistributions of source code must retain the above copyright
   9  *     notice, this list of conditions and the following disclaimer.
  10  * 2.  Redistributions in binary form must reproduce the above copyright
  11  *     notice, this list of conditions and the following disclaimer in the
  12  *     documentation and/or other materials provided with the distribution.
  13  * 3.  Neither the name of Apple Computer, Inc. ("Apple") nor the names of
  14  *     its contributors may be used to endorse or promote products derived
  15  *     from this software without specific prior written permission.
  16  *
  17  * THIS SOFTWARE IS PROVIDED BY APPLE AND ITS CONTRIBUTORS "AS IS" AND ANY
  18  * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  19  * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  20  * DISCLAIMED. IN NO EVENT SHALL APPLE OR ITS CONTRIBUTORS BE LIABLE FOR ANY
  21  * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  22  * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  23  * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  24  * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  25  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  26  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27  */
  28
  29 #ifndef RedBlackTree_h
  30 #define RedBlackTree_h
  31
  32 #include <wtf/Assertions.h>
  33 #include <wtf/Noncopyable.h>
  34
  35 namespace WTF {
  36
  37 // This implements a red-black tree with the following properties:
  38 // - The allocation of nodes in the tree is entirely up to the user.
  39 // - If you are in possession of a pointer to a node, you can delete
  40 //   it from the tree. The tree will subsequently no longer have a
  41 //   reference to this node.
  42 // - The key type must implement operator< and ==.
  43
  44 template<class NodeType, typename KeyType>
  45 class RedBlackTree {
  46     WTF_MAKE_NONCOPYABLE(RedBlackTree);
  47 private:
  48     enum Color {
  49         Red = 1,
  50         Black
  51     };
  52
  53 public:
  54     class Node {
  55         friend class RedBlackTree;
  56
  57     public:
  58         const NodeType* successor() const
  59         {
  60             const Node* x = this;
  61             if (x->right())
  62                 return treeMinimum(x->right());
  63             const NodeType* y = x->parent();
  64             while (y && x == y->right()) {
  65                 x = y;
  66                 y = y->parent();
  67             }
  68             return y;
  69         }
  70
  71         const NodeType* predecessor() const
  72         {
  73             const Node* x = this;
  74             if (x->left())
  75                 return treeMaximum(x->left());
  76             const NodeType* y = x->parent();
  77             while (y && x == y->left()) {
  78                 x = y;
  79                 y = y->parent();
  80             }
  81             return y;
  82         }
  83
  84         NodeType* successor()
  85         {
  86             return const_cast<NodeType*>(const_cast<const Node*>(this)->successor());
  87         }
  88
  89         NodeType* predecessor()
  90         {
  91             return const_cast<NodeType*>(const_cast<const Node*>(this)->predecessor());
  92         }
  93
  94     private:
  95         void reset()
  96         {
  97             m_left = 0;
  98             m_right = 0;
  99             m_parentAndRed = 1; // initialize to red
 100         }
 101
 102         // NOTE: these methods should pack the parent and red into a single
 103         // word. But doing so appears to reveal a bug in the compiler.
 104         NodeType* parent() const
 105         {
 106             return reinterpret_cast<NodeType*>(m_parentAndRed & ~static_cast<uintptr_t>(1));
 107         }
 108
 109         void setParent(NodeType* newParent)
 110         {
 111             m_parentAndRed = reinterpret_cast<uintptr_t>(newParent) | (m_parentAndRed & 1);
 112         }
 113
 114         NodeType* left() const
 115         {
 116             return m_left;
 117         }
 118
 119         void setLeft(NodeType* node)
 120         {
 121             m_left = node;
 122         }
 123
 124         NodeType* right() const
 125         {
 126             return m_right;
 127         }
 128
 129         void setRight(NodeType* node)
 130         {
 131             m_right = node;
 132         }
 133
 134         Color color() const
 135         {
 136             if (m_parentAndRed & 1)
 137                 return Red;
 138             return Black;
 139         }
 140
 141         void setColor(Color value)
 142         {
 143             if (value == Red)
 144                 m_parentAndRed |= 1;
 145             else
 146                 m_parentAndRed &= ~static_cast<uintptr_t>(1);
 147         }
 148
 149         NodeType* m_left;
 150         NodeType* m_right;
 151         uintptr_t m_parentAndRed;
 152     };
 153
 154     RedBlackTree()
 155         : m_root(0)
 156     {
 157     }
 158
 159     void insert(NodeType* x)
 160     {
 161         x->reset();
 162         treeInsert(x);
 163         x->setColor(Red);
 164
 165         while (x != m_root && x->parent()->color() == Red) {
 166             if (x->parent() == x->parent()->parent()->left()) {
 167                 NodeType* y = x->parent()->parent()->right();
 168                 if (y && y->color() == Red) {
 169                     // Case 1
 170                     x->parent()->setColor(Black);
 171                     y->setColor(Black);
 172                     x->parent()->parent()->setColor(Red);
 173                     x = x->parent()->parent();
 174                 } else {
 175                     if (x == x->parent()->right()) {
 176                         // Case 2
 177                         x = x->parent();
 178                         leftRotate(x);
 179                     }
 180                     // Case 3
 181                     x->parent()->setColor(Black);
 182                     x->parent()->parent()->setColor(Red);
 183                     rightRotate(x->parent()->parent());
 184                 }
 185             } else {
 186                 // Same as "then" clause with "right" and "left" exchanged.
 187                 NodeType* y = x->parent()->parent()->left();
 188                 if (y && y->color() == Red) {
 189                     // Case 1
 190                     x->parent()->setColor(Black);
 191                     y->setColor(Black);
 192                     x->parent()->parent()->setColor(Red);
 193                     x = x->parent()->parent();
 194                 } else {
 195                     if (x == x->parent()->left()) {
 196                         // Case 2
 197                         x = x->parent();
 198                         rightRotate(x);
 199                     }
 200                     // Case 3
 201                     x->parent()->setColor(Black);
 202                     x->parent()->parent()->setColor(Red);
 203                     leftRotate(x->parent()->parent());
 204                 }
 205             }
 206         }
 207
 208         m_root->setColor(Black);
 209     }
 210
 211     NodeType* remove(NodeType* z)
 212     {
 213         ASSERT(z);
 214         ASSERT(z->parent() || z == m_root);
 215
 216         // Y is the node to be unlinked from the tree.
 217         NodeType* y;
 218         if (!z->left() || !z->right())
 219             y = z;
 220         else
 221             y = z->successor();
 222
 223         // Y is guaranteed to be non-null at this point.
 224         NodeType* x;
 225         if (y->left())
 226             x = y->left();
 227         else
 228             x = y->right();
 229
 230         // X is the child of y which might potentially replace y in
 231         // the tree. X might be null at this point.
 232         NodeType* xParent;
 233         if (x) {
 234             x->setParent(y->parent());
 235             xParent = x->parent();
 236         } else
 237             xParent = y->parent();
 238         if (!y->parent())
 239             m_root = x;
 240         else {
 241             if (y == y->parent()->left())
 242                 y->parent()->setLeft(x);
 243             else
 244                 y->parent()->setRight(x);
 245         }
 246
 247         if (y != z) {
 248             if (y->color() == Black)
 249                 removeFixup(x, xParent);
 250
 251             y->setParent(z->parent());
 252             y->setColor(z->color());
 253             y->setLeft(z->left());
 254             y->setRight(z->right());
 255
 256             if (z->left())
 257                 z->left()->setParent(y);
 258             if (z->right())
 259                 z->right()->setParent(y);
 260             if (z->parent()) {
 261                 if (z->parent()->left() == z)
 262                     z->parent()->setLeft(y);
 263                 else
 264                     z->parent()->setRight(y);
 265             } else {
 266                 ASSERT(m_root == z);
 267                 m_root = y;
 268             }
 269         } else if (y->color() == Black)
 270             removeFixup(x, xParent);
 271
 272         return z;
 273     }
 274
 275     NodeType* remove(const KeyType& key)
 276     {
 277         NodeType* result = findExact(key);
 278         if (!result)
 279             return 0;
 280         return remove(result);
 281     }
 282
 283     NodeType* findExact(const KeyType& key) const
 284     {
 285         for (NodeType* current = m_root; current;) {
 286             if (current->key() == key)
 287                 return current;
 288             if (key < current->key())
 289                 current = current->left();
 290             else
 291                 current = current->right();
 292         }
 293         return 0;
 294     }
 295
 296     NodeType* findLeastGreaterThanOrEqual(const KeyType& key) const
 297     {
 298         NodeType* best = 0;
 299         for (NodeType* current = m_root; current;) {
 300             if (current->key() == key)
 301                 return current;
 302             if (current->key() < key)
 303                 current = current->right();
 304             else {
 305                 best = current;
 306                 current = current->left();
 307             }
 308         }
 309         return best;
 310     }
 311
 312     NodeType* findGreatestLessThanOrEqual(const KeyType& key) const
 313     {
 314         NodeType* best = 0;
 315         for (NodeType* current = m_root; current;) {
 316             if (current->key() == key)
 317                 return current;
 318             if (current->key() > key)
 319                 current = current->left();
 320             else {
 321                 best = current;
 322                 current = current->right();
 323             }
 324         }
 325         return best;
 326     }
 327
 328     NodeType* first() const
 329     {
 330         if (!m_root)
 331             return 0;
 332         return treeMinimum(m_root);
 333     }
 334
 335     NodeType* last() const
 336     {
 337         if (!m_root)
 338             return 0;
 339         return treeMaximum(m_root);
 340     }
 341
 342     // This is an O(n) operation.
 343     size_t size()
 344     {
 345         size_t result = 0;
 346         for (NodeType* current = first(); current; current = current->successor())
 347             result++;
 348         return result;
 349     }
 350
 351     // This is an O(1) operation.
 352     bool isEmpty()
 353     {
 354         return !m_root;
 355     }
 356
 357 private:
 358     // Finds the minimum element in the sub-tree rooted at the given
 359     // node.
 360     static NodeType* treeMinimum(NodeType* x)
 361     {
 362         while (x->left())
 363             x = x->left();
 364         return x;
 365     }
 366
 367     static NodeType* treeMaximum(NodeType* x)
 368     {
 369         while (x->right())
 370             x = x->right();
 371         return x;
 372     }
 373
 374     static const NodeType* treeMinimum(const NodeType* x)
 375     {
 376         while (x->left())
 377             x = x->left();
 378         return x;
 379     }
 380
 381     static const NodeType* treeMaximum(const NodeType* x)
 382     {
 383         while (x->right())
 384             x = x->right();
 385         return x;
 386     }
 387
 388     void treeInsert(NodeType* z)
 389     {
 390         ASSERT(!z->left());
 391         ASSERT(!z->right());
 392         ASSERT(!z->parent());
 393         ASSERT(z->color() == Red);
 394
 395         NodeType* y = 0;
 396         NodeType* x = m_root;
 397         while (x) {
 398             y = x;
 399             if (z->key() < x->key())
 400                 x = x->left();
 401             else
 402                 x = x->right();
 403         }
 404         z->setParent(y);
 405         if (!y)
 406             m_root = z;
 407         else {
 408             if (z->key() < y->key())
 409                 y->setLeft(z);
 410             else
 411                 y->setRight(z);
 412         }
 413     }
 414
 415     //----------------------------------------------------------------------
 416     // Red-Black tree operations
 417     //
 418
 419     // Left-rotates the subtree rooted at x.
 420     // Returns the new root of the subtree (x's right child).
 421     NodeType* leftRotate(NodeType* x)
 422     {
 423         // Set y.
 424         NodeType* y = x->right();
 425
 426         // Turn y's left subtree into x's right subtree.
 427         x->setRight(y->left());
 428         if (y->left())
 429             y->left()->setParent(x);
 430
 431         // Link x's parent to y.
 432         y->setParent(x->parent());
 433         if (!x->parent())
 434             m_root = y;
 435         else {
 436             if (x == x->parent()->left())
 437                 x->parent()->setLeft(y);
 438             else
 439                 x->parent()->setRight(y);
 440         }
 441
 442         // Put x on y's left.
 443         y->setLeft(x);
 444         x->setParent(y);
 445
 446         return y;
 447     }
 448
 449     // Right-rotates the subtree rooted at y.
 450     // Returns the new root of the subtree (y's left child).
 451     NodeType* rightRotate(NodeType* y)
 452     {
 453         // Set x.
 454         NodeType* x = y->left();
 455
 456         // Turn x's right subtree into y's left subtree.
 457         y->setLeft(x->right());
 458         if (x->right())
 459             x->right()->setParent(y);
 460
 461         // Link y's parent to x.
 462         x->setParent(y->parent());
 463         if (!y->parent())
 464             m_root = x;
 465         else {
 466             if (y == y->parent()->left())
 467                 y->parent()->setLeft(x);
 468             else
 469                 y->parent()->setRight(x);
 470         }
 471
 472         // Put y on x's right.
 473         x->setRight(y);
 474         y->setParent(x);
 475
 476         return x;
 477     }
 478
 479     // Restores the red-black property to the tree after splicing out
 480     // a node. Note that x may be null, which is why xParent must be
 481     // supplied.
 482     void removeFixup(NodeType* x, NodeType* xParent)
 483     {
 484         while (x != m_root && (!x || x->color() == Black)) {
 485             if (x == xParent->left()) {
 486                 // Note: the text points out that w can not be null.
 487                 // The reason is not obvious from simply looking at
 488                 // the code; it comes about from the properties of the
 489                 // red-black tree.
 490                 NodeType* w = xParent->right();
 491                 ASSERT(w); // x's sibling should not be null.
 492                 if (w->color() == Red) {
 493                     // Case 1
 494                     w->setColor(Black);
 495                     xParent->setColor(Red);
 496                     leftRotate(xParent);
 497                     w = xParent->right();
 498                 }
 499                 if ((!w->left() || w->left()->color() == Black)
 500                     && (!w->right() || w->right()->color() == Black)) {
 501                     // Case 2
 502                     w->setColor(Red);
 503                     x = xParent;
 504                     xParent = x->parent();
 505                 } else {
 506                     if (!w->right() || w->right()->color() == Black) {
 507                         // Case 3
 508                         w->left()->setColor(Black);
 509                         w->setColor(Red);
 510                         rightRotate(w);
 511                         w = xParent->right();
 512                     }
 513                     // Case 4
 514                     w->setColor(xParent->color());
 515                     xParent->setColor(Black);
 516                     if (w->right())
 517                         w->right()->setColor(Black);
 518                     leftRotate(xParent);
 519                     x = m_root;
 520                     xParent = x->parent();
 521                 }
 522             } else {
 523                 // Same as "then" clause with "right" and "left"
 524                 // exchanged.
 525
 526                 // Note: the text points out that w can not be null.
 527                 // The reason is not obvious from simply looking at
 528                 // the code; it comes about from the properties of the
 529                 // red-black tree.
 530                 NodeType* w = xParent->left();
 531                 ASSERT(w); // x's sibling should not be null.
 532                 if (w->color() == Red) {
 533                     // Case 1
 534                     w->setColor(Black);
 535                     xParent->setColor(Red);
 536                     rightRotate(xParent);
 537                     w = xParent->left();
 538                 }
 539                 if ((!w->right() || w->right()->color() == Black)
 540                     && (!w->left() || w->left()->color() == Black)) {
 541                     // Case 2
 542                     w->setColor(Red);
 543                     x = xParent;
 544                     xParent = x->parent();
 545                 } else {
 546                     if (!w->left() || w->left()->color() == Black) {
 547                         // Case 3
 548                         w->right()->setColor(Black);
 549                         w->setColor(Red);
 550                         leftRotate(w);
 551                         w = xParent->left();
 552                     }
 553                     // Case 4
 554                     w->setColor(xParent->color());
 555                     xParent->setColor(Black);
 556                     if (w->left())
 557                         w->left()->setColor(Black);
 558                     rightRotate(xParent);
 559                     x = m_root;
 560                     xParent = x->parent();
 561                 }
 562             }
 563         }
 564         if (x)
 565             x->setColor(Black);
 566     }
 567
 568     NodeType* m_root;
 569 };
 570
 571 }
 572
 573 #endif
 574