1/*
2 * Copyright (C) 2010 Google 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 *
14 * THIS SOFTWARE IS PROVIDED BY APPLE AND ITS CONTRIBUTORS "AS IS" AND ANY
15 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
16 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
17 * DISCLAIMED. IN NO EVENT SHALL APPLE OR ITS CONTRIBUTORS BE LIABLE FOR ANY
18 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
19 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
20 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
21 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
22 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
23 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
24 */
25
26// A red-black tree, which is a form of a balanced binary tree. It
27// supports efficient insertion, deletion and queries of comparable
28// elements. The same element may be inserted multiple times. The
29// algorithmic complexity of common operations is:
30//
31//   Insertion: O(lg(n))
32//   Deletion:  O(lg(n))
33//   Querying:  O(lg(n))
34//
35// The data type T that is stored in this red-black tree must be only
36// Plain Old Data (POD), or bottom out into POD. It must _not_ rely on
37// having its destructor called. This implementation internally
38// allocates storage in large chunks and does not call the destructor
39// on each object.
40//
41// Type T must supply a default constructor, a copy constructor, and
42// the "<" and "==" operators.
43//
44// In debug mode, printing of the data contained in the tree is
45// enabled. This requires the template specialization to be available:
46//
47//   template<> struct WebCore::ValueToString<T> {
48//       static String string(const T& t);
49//   };
50//
51// Note that when complex types are stored in this red/black tree, it
52// is possible that single invocations of the "<" and "==" operators
53// will be insufficient to describe the ordering of elements in the
54// tree during queries. As a concrete example, consider the case where
55// intervals are stored in the tree sorted by low endpoint. The "<"
56// operator on the Interval class only compares the low endpoint, but
57// the "==" operator takes into account the high endpoint as well.
58// This makes the necessary logic for querying and deletion somewhat
59// more complex. In order to properly handle such situations, the
60// property "needsFullOrderingComparisons" must be set to true on
61// the tree.
62//
63// This red-black tree is designed to be _augmented_; subclasses can
64// add additional and summary information to each node to efficiently
65// store and index more complex data structures. A concrete example is
66// the IntervalTree, which extends each node with a summary statistic
67// to efficiently store one-dimensional intervals.
68//
69// The design of this red-black tree comes from Cormen, Leiserson,
70// and Rivest, _Introduction to Algorithms_, MIT Press, 1990.
71
72#ifndef PODRedBlackTree_h
73#define PODRedBlackTree_h
74
75#include "PODArena.h"
76#include <wtf/Assertions.h>
77#include <wtf/Noncopyable.h>
78#include <wtf/RefPtr.h>
79#ifndef NDEBUG
80#include "Logging.h"
81#include <wtf/text/CString.h>
82#include <wtf/text/StringBuilder.h>
83#include <wtf/text/WTFString.h>
84#endif
85
86namespace WebCore {
87
88#ifndef NDEBUG
89template<class T>
90struct ValueToString;
91#endif
92
93template<class T>
94class PODRedBlackTree {
95public:
96    // Visitor interface for walking all of the tree's elements.
97    class Visitor {
98    public:
99        virtual void visit(const T& data) = 0;
100    protected:
101        virtual ~Visitor() { }
102    };
103
104    // Constructs a new red-black tree, allocating temporary objects
105    // from a newly constructed PODArena.
106    PODRedBlackTree()
107        : m_arena(PODArena::create())
108        , m_root(0)
109        , m_needsFullOrderingComparisons(false)
110#ifndef NDEBUG
111        , m_verboseDebugging(false)
112#endif
113    {
114    }
115
116    // Constructs a new red-black tree, allocating temporary objects
117    // from the given PODArena.
118    explicit PODRedBlackTree(PassRefPtr<PODArena> arena)
119        : m_arena(arena)
120        , m_root(0)
121        , m_needsFullOrderingComparisons(false)
122#ifndef NDEBUG
123        , m_verboseDebugging(false)
124#endif
125    {
126    }
127
128    virtual ~PODRedBlackTree() { }
129
130    void add(const T& data)
131    {
132        Node* node = m_arena->allocateObject<Node, T>(data);
133        insertNode(node);
134    }
135
136    // Returns true if the datum was found in the tree.
137    bool remove(const T& data)
138    {
139        Node* node = treeSearch(data);
140        if (node) {
141            deleteNode(node);
142            return true;
143        }
144        return false;
145    }
146
147    bool contains(const T& data) const { return treeSearch(data); }
148
149    void visitInorder(Visitor* visitor) const
150    {
151        if (!m_root)
152            return;
153        visitInorderImpl(m_root, visitor);
154    }
155
156    int size() const
157    {
158        Counter counter;
159        visitInorder(&counter);
160        return counter.count();
161    }
162
163    // See the class documentation for an explanation of this property.
164    void setNeedsFullOrderingComparisons(bool needsFullOrderingComparisons)
165    {
166        m_needsFullOrderingComparisons = needsFullOrderingComparisons;
167    }
168
169    virtual bool checkInvariants() const
170    {
171        int blackCount;
172        return checkInvariantsFromNode(m_root, &blackCount);
173    }
174
175#ifndef NDEBUG
176    // Dumps the tree's contents to the logging info stream for
177    // debugging purposes.
178    void dump() const
179    {
180        dumpFromNode(m_root, 0);
181    }
182
183    // Turns on or off verbose debugging of the tree, causing many
184    // messages to be logged during insertion and other operations in
185    // debug mode.
186    void setVerboseDebugging(bool verboseDebugging)
187    {
188        m_verboseDebugging = verboseDebugging;
189    }
190#endif
191
192protected:
193    enum Color {
194        Red = 1,
195        Black
196    };
197
198    // The base Node class which is stored in the tree. Nodes are only
199    // an internal concept; users of the tree deal only with the data
200    // they store in it.
201    class Node {
202        WTF_MAKE_NONCOPYABLE(Node);
203    public:
204        // Constructor. Newly-created nodes are colored red.
205        explicit Node(const T& data)
206            : m_left(0)
207            , m_right(0)
208            , m_parent(0)
209            , m_color(Red)
210            , m_data(data)
211        {
212        }
213
214        virtual ~Node() { }
215
216        Color color() const { return m_color; }
217        void setColor(Color color) { m_color = color; }
218
219        // Fetches the user data.
220        T& data() { return m_data; }
221
222        // Copies all user-level fields from the source node, but not
223        // internal fields. For example, the base implementation of this
224        // method copies the "m_data" field, but not the child or parent
225        // fields. Any augmentation information also does not need to be
226        // copied, as it will be recomputed. Subclasses must call the
227        // superclass implementation.
228        virtual void copyFrom(Node* src) { m_data = src->data(); }
229
230        Node* left() const { return m_left; }
231        void setLeft(Node* node) { m_left = node; }
232
233        Node* right() const { return m_right; }
234        void setRight(Node* node) { m_right = node; }
235
236        Node* parent() const { return m_parent; }
237        void setParent(Node* node) { m_parent = node; }
238
239    private:
240        Node* m_left;
241        Node* m_right;
242        Node* m_parent;
243        Color m_color;
244        T m_data;
245    };
246
247    // Returns the root of the tree, which is needed by some subclasses.
248    Node* root() const { return m_root; }
249
250private:
251    // This virtual method is the hook that subclasses should use when
252    // augmenting the red-black tree with additional per-node summary
253    // information. For example, in the case of an interval tree, this
254    // is used to compute the maximum endpoint of the subtree below the
255    // given node based on the values in the left and right children. It
256    // is guaranteed that this will be called in the correct order to
257    // properly update such summary information based only on the values
258    // in the left and right children. This method should return true if
259    // the node's summary information changed.
260    virtual bool updateNode(Node*) { return false; }
261
262    //----------------------------------------------------------------------
263    // Generic binary search tree operations
264    //
265
266    // Searches the tree for the given datum.
267    Node* treeSearch(const T& data) const
268    {
269        if (m_needsFullOrderingComparisons)
270            return treeSearchFullComparisons(m_root, data);
271
272        return treeSearchNormal(m_root, data);
273    }
274
275    // Searches the tree using the normal comparison operations,
276    // suitable for simple data types such as numbers.
277    Node* treeSearchNormal(Node* current, const T& data) const
278    {
279        while (current) {
280            if (current->data() == data)
281                return current;
282            if (data < current->data())
283                current = current->left();
284            else
285                current = current->right();
286        }
287        return 0;
288    }
289
290    // Searches the tree using multiple comparison operations, required
291    // for data types with more complex behavior such as intervals.
292    Node* treeSearchFullComparisons(Node* current, const T& data) const
293    {
294        if (!current)
295            return 0;
296        if (data < current->data())
297            return treeSearchFullComparisons(current->left(), data);
298        if (current->data() < data)
299            return treeSearchFullComparisons(current->right(), data);
300        if (data == current->data())
301            return current;
302
303        // We may need to traverse both the left and right subtrees.
304        Node* result = treeSearchFullComparisons(current->left(), data);
305        if (!result)
306            result = treeSearchFullComparisons(current->right(), data);
307        return result;
308    }
309
310    void treeInsert(Node* z)
311    {
312        Node* y = 0;
313        Node* x = m_root;
314        while (x) {
315            y = x;
316            if (z->data() < x->data())
317                x = x->left();
318            else
319                x = x->right();
320        }
321        z->setParent(y);
322        if (!y)
323            m_root = z;
324        else {
325            if (z->data() < y->data())
326                y->setLeft(z);
327            else
328                y->setRight(z);
329        }
330    }
331
332    // Finds the node following the given one in sequential ordering of
333    // their data, or null if none exists.
334    Node* treeSuccessor(Node* x)
335    {
336        if (x->right())
337            return treeMinimum(x->right());
338        Node* y = x->parent();
339        while (y && x == y->right()) {
340            x = y;
341            y = y->parent();
342        }
343        return y;
344    }
345
346    // Finds the minimum element in the sub-tree rooted at the given
347    // node.
348    Node* treeMinimum(Node* x)
349    {
350        while (x->left())
351            x = x->left();
352        return x;
353    }
354
355    // Helper for maintaining the augmented red-black tree.
356    void propagateUpdates(Node* start)
357    {
358        bool shouldContinue = true;
359        while (start && shouldContinue) {
360            shouldContinue = updateNode(start);
361            start = start->parent();
362        }
363    }
364
365    //----------------------------------------------------------------------
366    // Red-Black tree operations
367    //
368
369    // Left-rotates the subtree rooted at x.
370    // Returns the new root of the subtree (x's right child).
371    Node* leftRotate(Node* x)
372    {
373        // Set y.
374        Node* y = x->right();
375
376        // Turn y's left subtree into x's right subtree.
377        x->setRight(y->left());
378        if (y->left())
379            y->left()->setParent(x);
380
381        // Link x's parent to y.
382        y->setParent(x->parent());
383        if (!x->parent())
384            m_root = y;
385        else {
386            if (x == x->parent()->left())
387                x->parent()->setLeft(y);
388            else
389                x->parent()->setRight(y);
390        }
391
392        // Put x on y's left.
393        y->setLeft(x);
394        x->setParent(y);
395
396        // Update nodes lowest to highest.
397        updateNode(x);
398        updateNode(y);
399        return y;
400    }
401
402    // Right-rotates the subtree rooted at y.
403    // Returns the new root of the subtree (y's left child).
404    Node* rightRotate(Node* y)
405    {
406        // Set x.
407        Node* x = y->left();
408
409        // Turn x's right subtree into y's left subtree.
410        y->setLeft(x->right());
411        if (x->right())
412            x->right()->setParent(y);
413
414        // Link y's parent to x.
415        x->setParent(y->parent());
416        if (!y->parent())
417            m_root = x;
418        else {
419            if (y == y->parent()->left())
420                y->parent()->setLeft(x);
421            else
422                y->parent()->setRight(x);
423        }
424
425        // Put y on x's right.
426        x->setRight(y);
427        y->setParent(x);
428
429        // Update nodes lowest to highest.
430        updateNode(y);
431        updateNode(x);
432        return x;
433    }
434
435    // Inserts the given node into the tree.
436    void insertNode(Node* x)
437    {
438        treeInsert(x);
439        x->setColor(Red);
440        updateNode(x);
441
442        logIfVerbose("  PODRedBlackTree::InsertNode");
443
444        // The node from which to start propagating updates upwards.
445        Node* updateStart = x->parent();
446
447        while (x != m_root && x->parent()->color() == Red) {
448            if (x->parent() == x->parent()->parent()->left()) {
449                Node* y = x->parent()->parent()->right();
450                if (y && y->color() == Red) {
451                    // Case 1
452                    logIfVerbose("  Case 1/1");
453                    x->parent()->setColor(Black);
454                    y->setColor(Black);
455                    x->parent()->parent()->setColor(Red);
456                    updateNode(x->parent());
457                    x = x->parent()->parent();
458                    updateNode(x);
459                    updateStart = x->parent();
460                } else {
461                    if (x == x->parent()->right()) {
462                        logIfVerbose("  Case 1/2");
463                        // Case 2
464                        x = x->parent();
465                        leftRotate(x);
466                    }
467                    // Case 3
468                    logIfVerbose("  Case 1/3");
469                    x->parent()->setColor(Black);
470                    x->parent()->parent()->setColor(Red);
471                    Node* newSubTreeRoot = rightRotate(x->parent()->parent());
472                    updateStart = newSubTreeRoot->parent();
473                }
474            } else {
475                // Same as "then" clause with "right" and "left" exchanged.
476                Node* y = x->parent()->parent()->left();
477                if (y && y->color() == Red) {
478                    // Case 1
479                    logIfVerbose("  Case 2/1");
480                    x->parent()->setColor(Black);
481                    y->setColor(Black);
482                    x->parent()->parent()->setColor(Red);
483                    updateNode(x->parent());
484                    x = x->parent()->parent();
485                    updateNode(x);
486                    updateStart = x->parent();
487                } else {
488                    if (x == x->parent()->left()) {
489                        // Case 2
490                        logIfVerbose("  Case 2/2");
491                        x = x->parent();
492                        rightRotate(x);
493                    }
494                    // Case 3
495                    logIfVerbose("  Case 2/3");
496                    x->parent()->setColor(Black);
497                    x->parent()->parent()->setColor(Red);
498                    Node* newSubTreeRoot = leftRotate(x->parent()->parent());
499                    updateStart = newSubTreeRoot->parent();
500                }
501            }
502        }
503
504        propagateUpdates(updateStart);
505
506        m_root->setColor(Black);
507    }
508
509    // Restores the red-black property to the tree after splicing out
510    // a node. Note that x may be null, which is why xParent must be
511    // supplied.
512    void deleteFixup(Node* x, Node* xParent)
513    {
514        while (x != m_root && (!x || x->color() == Black)) {
515            if (x == xParent->left()) {
516                // Note: the text points out that w can not be null.
517                // The reason is not obvious from simply looking at
518                // the code; it comes about from the properties of the
519                // red-black tree.
520                Node* w = xParent->right();
521                ASSERT(w); // x's sibling should not be null.
522                if (w->color() == Red) {
523                    // Case 1
524                    w->setColor(Black);
525                    xParent->setColor(Red);
526                    leftRotate(xParent);
527                    w = xParent->right();
528                }
529                if ((!w->left() || w->left()->color() == Black)
530                    && (!w->right() || w->right()->color() == Black)) {
531                    // Case 2
532                    w->setColor(Red);
533                    x = xParent;
534                    xParent = x->parent();
535                } else {
536                    if (!w->right() || w->right()->color() == Black) {
537                        // Case 3
538                        w->left()->setColor(Black);
539                        w->setColor(Red);
540                        rightRotate(w);
541                        w = xParent->right();
542                    }
543                    // Case 4
544                    w->setColor(xParent->color());
545                    xParent->setColor(Black);
546                    if (w->right())
547                        w->right()->setColor(Black);
548                    leftRotate(xParent);
549                    x = m_root;
550                    xParent = x->parent();
551                }
552            } else {
553                // Same as "then" clause with "right" and "left"
554                // exchanged.
555
556                // Note: the text points out that w can not be null.
557                // The reason is not obvious from simply looking at
558                // the code; it comes about from the properties of the
559                // red-black tree.
560                Node* w = xParent->left();
561                ASSERT(w); // x's sibling should not be null.
562                if (w->color() == Red) {
563                    // Case 1
564                    w->setColor(Black);
565                    xParent->setColor(Red);
566                    rightRotate(xParent);
567                    w = xParent->left();
568                }
569                if ((!w->right() || w->right()->color() == Black)
570                    && (!w->left() || w->left()->color() == Black)) {
571                    // Case 2
572                    w->setColor(Red);
573                    x = xParent;
574                    xParent = x->parent();
575                } else {
576                    if (!w->left() || w->left()->color() == Black) {
577                        // Case 3
578                        w->right()->setColor(Black);
579                        w->setColor(Red);
580                        leftRotate(w);
581                        w = xParent->left();
582                    }
583                    // Case 4
584                    w->setColor(xParent->color());
585                    xParent->setColor(Black);
586                    if (w->left())
587                        w->left()->setColor(Black);
588                    rightRotate(xParent);
589                    x = m_root;
590                    xParent = x->parent();
591                }
592            }
593        }
594        if (x)
595            x->setColor(Black);
596    }
597
598    // Deletes the given node from the tree. Note that this
599    // particular node may not actually be removed from the tree;
600    // instead, another node might be removed and its contents
601    // copied into z.
602    void deleteNode(Node* z)
603    {
604        // Y is the node to be unlinked from the tree.
605        Node* y;
606        if (!z->left() || !z->right())
607            y = z;
608        else
609            y = treeSuccessor(z);
610
611        // Y is guaranteed to be non-null at this point.
612        Node* x;
613        if (y->left())
614            x = y->left();
615        else
616            x = y->right();
617
618        // X is the child of y which might potentially replace y in
619        // the tree. X might be null at this point.
620        Node* xParent;
621        if (x) {
622            x->setParent(y->parent());
623            xParent = x->parent();
624        } else
625            xParent = y->parent();
626        if (!y->parent())
627            m_root = x;
628        else {
629            if (y == y->parent()->left())
630                y->parent()->setLeft(x);
631            else
632                y->parent()->setRight(x);
633        }
634        if (y != z) {
635            z->copyFrom(y);
636            // This node has changed location in the tree and must be updated.
637            updateNode(z);
638            // The parent and its parents may now be out of date.
639            propagateUpdates(z->parent());
640        }
641
642        // If we haven't already updated starting from xParent, do so now.
643        if (xParent && xParent != y && xParent != z)
644            propagateUpdates(xParent);
645        if (y->color() == Black)
646            deleteFixup(x, xParent);
647    }
648
649    // Visits the subtree rooted at the given node in order.
650    void visitInorderImpl(Node* node, Visitor* visitor) const
651    {
652        if (node->left())
653            visitInorderImpl(node->left(), visitor);
654        visitor->visit(node->data());
655        if (node->right())
656            visitInorderImpl(node->right(), visitor);
657    }
658
659    //----------------------------------------------------------------------
660    // Helper class for size()
661
662    // A Visitor which simply counts the number of visited elements.
663    class Counter : public Visitor {
664        WTF_MAKE_NONCOPYABLE(Counter);
665    public:
666        Counter()
667            : m_count(0) { }
668
669        virtual void visit(const T& data) { ++m_count; }
670        int count() const { return m_count; }
671
672    private:
673        int m_count;
674    };
675
676    //----------------------------------------------------------------------
677    // Verification and debugging routines
678    //
679
680    // Returns in the "blackCount" parameter the number of black
681    // children along all paths to all leaves of the given node.
682    bool checkInvariantsFromNode(Node* node, int* blackCount) const
683    {
684        // Base case is a leaf node.
685        if (!node) {
686            *blackCount = 1;
687            return true;
688        }
689
690        // Each node is either red or black.
691        if (!(node->color() == Red || node->color() == Black))
692            return false;
693
694        // Every leaf (or null) is black.
695
696        if (node->color() == Red) {
697            // Both of its children are black.
698            if (!((!node->left() || node->left()->color() == Black)))
699                return false;
700            if (!((!node->right() || node->right()->color() == Black)))
701                return false;
702        }
703
704        // Every simple path to a leaf node contains the same number of
705        // black nodes.
706        int leftCount = 0, rightCount = 0;
707        bool leftValid = checkInvariantsFromNode(node->left(), &leftCount);
708        bool rightValid = checkInvariantsFromNode(node->right(), &rightCount);
709        if (!leftValid || !rightValid)
710            return false;
711        *blackCount = leftCount + (node->color() == Black ? 1 : 0);
712        return leftCount == rightCount;
713    }
714
715#ifdef NDEBUG
716    void logIfVerbose(const char*) const { }
717#else
718    void logIfVerbose(const char* output) const
719    {
720        if (m_verboseDebugging)
721            LOG_ERROR("%s", output);
722    }
723#endif
724
725#ifndef NDEBUG
726    // Dumps the subtree rooted at the given node.
727    void dumpFromNode(Node* node, int indentation) const
728    {
729        StringBuilder builder;
730        for (int i = 0; i < indentation; i++)
731            builder.append(" ");
732        builder.append("-");
733        if (node) {
734            builder.append(" ");
735            builder.append(ValueToString<T>::string(node->data()));
736            builder.append((node->color() == Black) ? " (black)" : " (red)");
737        }
738        LOG_ERROR("%s", builder.toString().ascii().data());
739        if (node) {
740            dumpFromNode(node->left(), indentation + 2);
741            dumpFromNode(node->right(), indentation + 2);
742        }
743    }
744#endif
745
746    //----------------------------------------------------------------------
747    // Data members
748
749    RefPtr<PODArena> m_arena;
750    Node* m_root;
751    bool m_needsFullOrderingComparisons;
752#ifndef NDEBUG
753    bool m_verboseDebugging;
754#endif
755};
756
757} // namespace WebCore
758
759#endif // PODRedBlackTree_h
760