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 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#pragma once
30
31#include <wtf/Assertions.h>
32#include <wtf/Noncopyable.h>
33
34namespace WTF {
35
36// This implements a red-black tree with the following properties:
37// - The allocation of nodes in the tree is entirely up to the user.
38// - If you are in possession of a pointer to a node, you can delete
39// it from the tree. The tree will subsequently no longer have a
40// reference to this node.
41// - The key type must implement operator< and ==.
42
43template<class NodeType, typename KeyType>
44class RedBlackTree final {
45 WTF_MAKE_NONCOPYABLE(RedBlackTree);
46 WTF_MAKE_FAST_ALLOCATED;
47private:
48 enum Color {
49 Red = 1,
50 Black
51 };
52
53public:
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
357private:
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