1 // util/rbtree.h -- A red/black tree implementation
3 // This software is copyright (c) 2006 Scott Wood <scott@buserror.net>.
5 // This software is provided 'as-is', without any express or implied warranty.
6 // In no event will the authors or contributors be held liable for any damages
7 // arising from the use of this software.
9 // Permission is hereby granted to everyone, free of charge, to use, copy,
10 // modify, prepare derivative works of, publish, distribute, perform,
11 // sublicense, and/or sell copies of the Software, provided that the above
12 // copyright notice and disclaimer of warranty be included in all copies or
13 // substantial portions of this software.
15 #ifndef _UTIL_RBTREE_H
16 #define _UTIL_RBTREE_H
23 // T must have an RBTree<T, NodeV, KeyV>::Node member called rbtree_node.
24 // NodeV < NodeV, NodeV < KeyV, and NodeV > KeyV must be supported.
26 // Using NodeV != KeyV allows things such as having NodeV represent a
27 // range of values and KeyV a single value, causing find() to search
28 // for the node in whose range KeyV falls. It is an error to add
29 // nodes that are not well ordered.
31 template <typename T, typename NodeV, typename KeyV>
35 Node *parent, *left, *right;
39 // A pointer to the parent's left or right pointer
40 // (whichever side this node is on). This is also
41 // used for sanity checks (it is NULL for nodes not
48 parent = left = right = NULL;
54 return parentlink != NULL;
61 T *node_to_type(Node *n)
63 return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(n) -
64 offsetof(T, rbtree_node));
67 void assert_rb_condition(Node *left, Node *right)
69 assert(left->value < right->value);
70 assert(!(right->value < left->value));
71 assert(!left->red || !right->red);
74 void assert_well_ordered(Node *left, NodeV right)
76 assert(left->value < right);
77 assert(!(left->value > right));
78 assert(!left->red || !right->red);
81 Node *find_node(Node *n, NodeV &val, Node ***link);
82 Node *find_node_key(KeyV val, bool exact = true);
84 bool node_is_left(Node *n)
86 assert(n->parentlink == &n->parent->left ||
87 n->parentlink == &n->parent->right);
89 return n->parentlink == &n->parent->left;
92 Node *sibling(Node *n)
95 return n->parent->right;
97 return n->parent->left;
102 return sibling(n->parent);
105 void rotate_left(Node *n);
106 void rotate_right(Node *n);
108 // B must have no right child and have A as an ancestor.
109 void swap(Node *a, Node *b);
129 Node *n = find_node_key(value);
132 return node_to_type(n);
137 T *find_nearest(KeyV value)
139 Node *n = find_node_key(value, false);
142 return node_to_type(n);
144 // Should only happen on an empty tree
152 template <typename T, typename NodeV, typename KeyV>
153 typename RBTree<T, NodeV, KeyV>::Node *
154 RBTree<T, NodeV, KeyV>::find_node(Node *n, NodeV &val, Node ***link)
160 *link = n->parentlink;
170 if (n->value < val) {
174 assert_rb_condition(n, next);
182 // This assert detects duplicate nodes, but not
183 // overlapping ranges (or otherwise non-well-ordered
186 assert(val < n->value);
190 assert_rb_condition(next, n);
205 template <typename T, typename NodeV, typename KeyV>
206 typename RBTree<T, NodeV, KeyV>::Node *
207 RBTree<T, NodeV, KeyV>::find_node_key(KeyV val, bool exact)
214 if (n->value < val) {
218 assert_rb_condition(n, next);
219 } else if (n->value > val) {
223 assert_rb_condition(next, n);
237 template <typename T, typename NodeV, typename KeyV>
238 void RBTree<T, NodeV, KeyV>::rotate_left(Node *n)
240 Node *new_top = n->right;
242 assert(*n->parentlink == n);
243 *n->parentlink = new_top;
245 assert(new_top->parent == n);
246 assert(new_top->parentlink == &n->right);
248 new_top->parent = n->parent;
249 new_top->parentlink = n->parentlink;
252 n->parentlink = &new_top->left;
254 n->right = new_top->left;
258 assert(n->right->parent == new_top);
259 assert(n->right->parentlink == &new_top->left);
261 n->right->parent = n;
262 n->right->parentlink = &n->right;
266 template <typename T, typename NodeV, typename KeyV>
267 void RBTree<T, NodeV, KeyV>::rotate_right(Node *n)
269 Node *new_top = n->left;
271 assert(*n->parentlink == n);
272 *n->parentlink = new_top;
274 assert(new_top->parent == n);
275 assert(new_top->parentlink == &n->left);
277 new_top->parent = n->parent;
278 new_top->parentlink = n->parentlink;
281 n->parentlink = &new_top->right;
283 n->left = new_top->right;
287 assert(n->left->parent == new_top);
288 assert(n->left->parentlink == &new_top->right);
291 n->left->parentlink = &n->left;
295 // B must have no right child and have A as an ancestor.
296 template <typename T, typename NodeV, typename KeyV>
297 void RBTree<T, NodeV, KeyV>::swap(Node *a, Node *b)
299 Node *bparent = b->parent;
300 Node **bparentlink = b->parentlink;
303 assert(a->left || a->right);
305 b->parent = a->parent;
306 b->parentlink = a->parentlink;
310 a->parentlink = &b->left;
313 a->parentlink = bparentlink;
316 assert(a->parent != a);
317 assert(b->parent != b);
319 Node *bleft = b->left;
329 assert(a->parent != a);
330 assert(b->parent != b);
338 a->left->parentlink = &a->left;
342 b->right->parent = b;
343 b->right->parentlink = &b->right;
348 b->left->parentlink = &b->left;
351 template <typename T, typename NodeV, typename KeyV>
352 void RBTree<T, NodeV, KeyV>::add(T *t)
354 Node *n = &t->rbtree_node;
355 assert(!n->is_on_rbtree());
357 Node *insert_at = find_node(top, n->value, &n->parentlink);
359 assert(insert_at || !top);
361 n->parent = insert_at;
362 n->left = n->right = NULL;
365 assert(n->parentlink);
366 assert(*n->parentlink == n);
373 assert((n->parent->value < n->value) !=
374 (n->value < n->parent->value));
380 Node *unc = uncle(n);
382 assert(!n->parent->parent->red);
383 n->parent->red = unc->red = false;
384 n = n->parent->parent;
388 if (node_is_left(n)) {
389 if (!node_is_left(n->parent)) {
390 rotate_right(n->parent);
394 if (node_is_left(n->parent)) {
395 rotate_left(n->parent);
400 assert(n->parent->red);
401 assert(!red(uncle(n)));
402 assert(!n->parent->parent->red);
404 n->parent->red = false;
405 n->parent->parent->red = true;
407 if (node_is_left(n)) {
408 assert(node_is_left(n->parent));
409 rotate_right(n->parent->parent);
411 assert(!node_is_left(n->parent));
412 rotate_left(n->parent->parent);
416 template <typename T, typename NodeV, typename KeyV>
417 void RBTree<T, NodeV, KeyV>::del(T *t)
419 Node *n = &t->rbtree_node;
420 assert(*n->parentlink == n);
422 if (n->left && n->right) {
423 Node *highest_on_left = find_node(n->left, n->value, NULL);
424 assert(!highest_on_left->right);
425 swap(n, highest_on_left);
429 Node *parent = n->parent;
430 Node *child = n->left ? n->left : n->right;
433 assert(child->parent == n);
435 child->parent = n->parent;
436 child->parentlink = n->parentlink;
437 *child->parentlink = child;
438 assert(child != parent);
440 *n->parentlink = NULL;
443 n->parentlink = NULL;
466 sib = parent->left ? parent->left : parent->right;
469 assert(!parent->red);
470 assert(!red(sib->left));
471 assert(!red(sib->right));
476 if (node_is_left(sib)) {
477 rotate_right(parent);
485 assert(sib == sibling(n));
486 } else if (!parent->red && !red(sib->left) && !red(sib->right)) {
490 parent = parent->parent;
496 if (!parent->red && !red(sib->left) && !red(sib->right)) {
502 if (!red(sib->left) && !red(sib->right)) {
509 if (node_is_left(sib) && !red(sib->left) && red(sib->right)) {
511 sib->right->red = false;
514 } else if (!node_is_left(sib) && !red(sib->right) && red(sib->left)) {
516 sib->left->red = false;
521 assert(parent == sib->parent);
524 sib->red = parent->red;
527 if (node_is_left(sib)) {
528 assert(sib->left->red);
529 sib->left->red = false;
530 rotate_right(parent);
532 assert(sib->right->red);
533 sib->right->red = false;
538 // RBPtr is a Pointer->Value associative array, and RBInt is an
539 // Integer->Value associative array.
541 template<typename Ptr, typename Val>
543 typedef RBTree<RBPtrNode, Ptr, Ptr> Tree;
544 typename Tree::Node rbtree_node;
547 intptr_t operator < (RBPtrNode &other)
549 return (intptr_t)other.rbtree_node.value -
550 (intptr_t)rbtree_node.value;
553 intptr_t operator > (RBPtrNode &other)
555 return (intptr_t)rbtree_node.value -
556 (intptr_t)other.rbtree_node.value;
565 template<typename Ptr, typename Val>
566 struct RBPtr : public RBTree<RBPtrNode<Ptr, Val>, Ptr, Ptr>
568 typedef RBPtrNode<Ptr, Val> Node;
569 typedef RBTree<Node, Ptr, Ptr> Tree;
571 void add(Ptr ptr, Val &val)
573 Node *node = new Node;
575 node->rbtree_node.value = ptr;
585 template<typename Int, typename Val>
587 typedef RBTree<RBIntNode, Int, Int> Tree;
588 typename Tree::Node rbtree_node;
591 intptr_t operator < (RBIntNode &other)
593 return other.rbtree_node.value - rbtree_node.value;
596 intptr_t operator > (RBIntNode &other)
598 return rbtree_node.value - other.rbtree_node.value;
607 template<typename Int, typename Val>
608 struct RBInt : public RBTree<RBIntNode<Int, Val>, Int, Int>
610 typedef RBIntNode<Int, Val> Node;
611 typedef RBTree<Node, Int, Int> Tree;
613 void add(Int key, Val &val)
615 Node *node = new Node;
617 node->rbtree_node.value = key;