3 (C) 1999 Andrea Arcangeli <andrea@suse.de>
4 (C) 2002 David Woodhouse <dwmw2@infradead.org>
5 (C) 2012 Michel Lespinasse <walken@google.com>
7 This program is free software; you can redistribute it and/or modify
8 it under the terms of the GNU General Public License as published by
9 the Free Software Foundation; either version 2 of the License, or
10 (at your option) any later version.
12 This program is distributed in the hope that it will be useful,
13 but WITHOUT ANY WARRANTY; without even the implied warranty of
14 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15 GNU General Public License for more details.
17 You should have received a copy of the GNU General Public License
18 along with this program; if not, write to the Free Software
19 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
24 #include <linux/rbtree.h>
25 #include <linux/export.h>
28 * red-black trees properties: http://en.wikipedia.org/wiki/Rbtree
30 * 1) A node is either red or black
31 * 2) The root is black
32 * 3) All leaves (NULL) are black
33 * 4) Both children of every red node are black
34 * 5) Every simple path from root to leaves contains the same number
37 * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
38 * consecutive red nodes in a path and every red node is therefore followed by
39 * a black. So if B is the number of black nodes on every simple path (as per
40 * 5), then the longest possible path due to 4 is 2B.
42 * We shall indicate color with case, where black nodes are uppercase and red
43 * nodes will be lowercase. Unknown color nodes shall be drawn as red within
44 * parentheses and have some accompanying text comment.
50 #define __rb_parent(pc) ((struct rb_node *)(pc & ~3))
52 #define __rb_color(pc) ((pc) & 1)
53 #define __rb_is_black(pc) __rb_color(pc)
54 #define __rb_is_red(pc) (!__rb_color(pc))
55 #define rb_color(rb) __rb_color((rb)->__rb_parent_color)
56 #define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
57 #define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)
59 static inline void rb_set_black(struct rb_node *rb)
61 rb->__rb_parent_color |= RB_BLACK;
64 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
66 rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
69 static inline void rb_set_parent_color(struct rb_node *rb,
70 struct rb_node *p, int color)
72 rb->__rb_parent_color = (unsigned long)p | color;
75 static inline struct rb_node *rb_red_parent(struct rb_node *red)
77 return (struct rb_node *)red->__rb_parent_color;
81 __rb_change_child(struct rb_node *old, struct rb_node *new,
82 struct rb_node *parent, struct rb_root *root)
85 if (parent->rb_left == old)
86 parent->rb_left = new;
88 parent->rb_right = new;
94 * Helper function for rotations:
95 * - old's parent and color get assigned to new
96 * - old gets assigned new as a parent and 'color' as a color.
99 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
100 struct rb_root *root, int color)
102 struct rb_node *parent = rb_parent(old);
103 new->__rb_parent_color = old->__rb_parent_color;
104 rb_set_parent_color(old, new, color);
105 __rb_change_child(old, new, parent, root);
108 static __always_inline void
109 __rb_insert(struct rb_node *node, struct rb_root *root,
110 void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
112 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
116 * Loop invariant: node is red
118 * If there is a black parent, we are done.
119 * Otherwise, take some corrective action as we don't
120 * want a red root or two consecutive red nodes.
123 rb_set_parent_color(node, NULL, RB_BLACK);
125 } else if (rb_is_black(parent))
128 gparent = rb_red_parent(parent);
130 tmp = gparent->rb_right;
131 if (parent != tmp) { /* parent == gparent->rb_left */
132 if (tmp && rb_is_red(tmp)) {
134 * Case 1 - color flips
142 * However, since g's parent might be red, and
143 * 4) does not allow this, we need to recurse
146 rb_set_parent_color(tmp, gparent, RB_BLACK);
147 rb_set_parent_color(parent, gparent, RB_BLACK);
149 parent = rb_parent(node);
150 rb_set_parent_color(node, parent, RB_RED);
154 tmp = parent->rb_right;
157 * Case 2 - left rotate at parent
165 * This still leaves us in violation of 4), the
166 * continuation into Case 3 will fix that.
168 parent->rb_right = tmp = node->rb_left;
169 node->rb_left = parent;
171 rb_set_parent_color(tmp, parent,
173 rb_set_parent_color(parent, node, RB_RED);
174 augment_rotate(parent, node);
176 tmp = node->rb_right;
180 * Case 3 - right rotate at gparent
188 gparent->rb_left = tmp; /* == parent->rb_right */
189 parent->rb_right = gparent;
191 rb_set_parent_color(tmp, gparent, RB_BLACK);
192 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
193 augment_rotate(gparent, parent);
196 tmp = gparent->rb_left;
197 if (tmp && rb_is_red(tmp)) {
198 /* Case 1 - color flips */
199 rb_set_parent_color(tmp, gparent, RB_BLACK);
200 rb_set_parent_color(parent, gparent, RB_BLACK);
202 parent = rb_parent(node);
203 rb_set_parent_color(node, parent, RB_RED);
207 tmp = parent->rb_left;
209 /* Case 2 - right rotate at parent */
210 parent->rb_left = tmp = node->rb_right;
211 node->rb_right = parent;
213 rb_set_parent_color(tmp, parent,
215 rb_set_parent_color(parent, node, RB_RED);
216 augment_rotate(parent, node);
221 /* Case 3 - left rotate at gparent */
222 gparent->rb_right = tmp; /* == parent->rb_left */
223 parent->rb_left = gparent;
225 rb_set_parent_color(tmp, gparent, RB_BLACK);
226 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
227 augment_rotate(gparent, parent);
233 static __always_inline void
234 __rb_erase_color(struct rb_node *parent, struct rb_root *root,
235 const struct rb_augment_callbacks *augment)
237 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
242 * - node is black (or NULL on first iteration)
243 * - node is not the root (parent is not NULL)
244 * - All leaf paths going through parent and node have a
245 * black node count that is 1 lower than other leaf paths.
247 sibling = parent->rb_right;
248 if (node != sibling) { /* node == parent->rb_left */
249 if (rb_is_red(sibling)) {
251 * Case 1 - left rotate at parent
259 parent->rb_right = tmp1 = sibling->rb_left;
260 sibling->rb_left = parent;
261 rb_set_parent_color(tmp1, parent, RB_BLACK);
262 __rb_rotate_set_parents(parent, sibling, root,
264 augment->rotate(parent, sibling);
267 tmp1 = sibling->rb_right;
268 if (!tmp1 || rb_is_black(tmp1)) {
269 tmp2 = sibling->rb_left;
270 if (!tmp2 || rb_is_black(tmp2)) {
272 * Case 2 - sibling color flip
273 * (p could be either color here)
281 * This leaves us violating 5) which
282 * can be fixed by flipping p to black
283 * if it was red, or by recursing at p.
284 * p is red when coming from Case 1.
286 rb_set_parent_color(sibling, parent,
288 if (rb_is_red(parent))
289 rb_set_black(parent);
292 parent = rb_parent(node);
299 * Case 3 - right rotate at sibling
300 * (p could be either color here)
310 sibling->rb_left = tmp1 = tmp2->rb_right;
311 tmp2->rb_right = sibling;
312 parent->rb_right = tmp2;
314 rb_set_parent_color(tmp1, sibling,
316 augment->rotate(sibling, tmp2);
321 * Case 4 - left rotate at parent + color flips
322 * (p and sl could be either color here.
323 * After rotation, p becomes black, s acquires
324 * p's color, and sl keeps its color)
332 parent->rb_right = tmp2 = sibling->rb_left;
333 sibling->rb_left = parent;
334 rb_set_parent_color(tmp1, sibling, RB_BLACK);
336 rb_set_parent(tmp2, parent);
337 __rb_rotate_set_parents(parent, sibling, root,
339 augment->rotate(parent, sibling);
342 sibling = parent->rb_left;
343 if (rb_is_red(sibling)) {
344 /* Case 1 - right rotate at parent */
345 parent->rb_left = tmp1 = sibling->rb_right;
346 sibling->rb_right = parent;
347 rb_set_parent_color(tmp1, parent, RB_BLACK);
348 __rb_rotate_set_parents(parent, sibling, root,
350 augment->rotate(parent, sibling);
353 tmp1 = sibling->rb_left;
354 if (!tmp1 || rb_is_black(tmp1)) {
355 tmp2 = sibling->rb_right;
356 if (!tmp2 || rb_is_black(tmp2)) {
357 /* Case 2 - sibling color flip */
358 rb_set_parent_color(sibling, parent,
360 if (rb_is_red(parent))
361 rb_set_black(parent);
364 parent = rb_parent(node);
370 /* Case 3 - right rotate at sibling */
371 sibling->rb_right = tmp1 = tmp2->rb_left;
372 tmp2->rb_left = sibling;
373 parent->rb_left = tmp2;
375 rb_set_parent_color(tmp1, sibling,
377 augment->rotate(sibling, tmp2);
381 /* Case 4 - left rotate at parent + color flips */
382 parent->rb_left = tmp2 = sibling->rb_right;
383 sibling->rb_right = parent;
384 rb_set_parent_color(tmp1, sibling, RB_BLACK);
386 rb_set_parent(tmp2, parent);
387 __rb_rotate_set_parents(parent, sibling, root,
389 augment->rotate(parent, sibling);
395 static __always_inline void
396 __rb_erase(struct rb_node *node, struct rb_root *root,
397 const struct rb_augment_callbacks *augment)
399 struct rb_node *child = node->rb_right, *tmp = node->rb_left;
400 struct rb_node *parent, *rebalance;
405 * Case 1: node to erase has no more than 1 child (easy!)
407 * Note that if there is one child it must be red due to 5)
408 * and node must be black due to 4). We adjust colors locally
409 * so as to bypass __rb_erase_color() later on.
411 pc = node->__rb_parent_color;
412 parent = __rb_parent(pc);
413 __rb_change_child(node, child, parent, root);
415 child->__rb_parent_color = pc;
418 rebalance = __rb_is_black(pc) ? parent : NULL;
421 /* Still case 1, but this time the child is node->rb_left */
422 tmp->__rb_parent_color = pc = node->__rb_parent_color;
423 parent = __rb_parent(pc);
424 __rb_change_child(node, tmp, parent, root);
428 struct rb_node *successor = child, *child2;
429 tmp = child->rb_left;
432 * Case 2: node's successor is its right child
441 child2 = successor->rb_right;
442 augment->copy(node, successor);
445 * Case 3: node's successor is leftmost under
446 * node's right child subtree
463 parent->rb_left = child2 = successor->rb_right;
464 successor->rb_right = child;
465 rb_set_parent(child, successor);
466 augment->copy(node, successor);
467 augment->propagate(parent, successor);
470 successor->rb_left = tmp = node->rb_left;
471 rb_set_parent(tmp, successor);
473 pc = node->__rb_parent_color;
474 tmp = __rb_parent(pc);
475 __rb_change_child(node, successor, tmp, root);
477 successor->__rb_parent_color = pc;
478 rb_set_parent_color(child2, parent, RB_BLACK);
481 unsigned long pc2 = successor->__rb_parent_color;
482 successor->__rb_parent_color = pc;
483 rebalance = __rb_is_black(pc2) ? parent : NULL;
488 augment->propagate(tmp, NULL);
490 __rb_erase_color(rebalance, root, augment);
494 * Non-augmented rbtree manipulation functions.
496 * We use dummy augmented callbacks here, and have the compiler optimize them
497 * out of the rb_insert_color() and rb_erase() function definitions.
500 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
501 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
502 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
504 static const struct rb_augment_callbacks dummy_callbacks = {
505 dummy_propagate, dummy_copy, dummy_rotate
508 void rb_insert_color(struct rb_node *node, struct rb_root *root)
510 __rb_insert(node, root, dummy_rotate);
512 EXPORT_SYMBOL(rb_insert_color);
514 void rb_erase(struct rb_node *node, struct rb_root *root)
516 __rb_erase(node, root, &dummy_callbacks);
518 EXPORT_SYMBOL(rb_erase);
521 * Augmented rbtree manipulation functions.
523 * This instantiates the same __always_inline functions as in the non-augmented
524 * case, but this time with user-defined callbacks.
527 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
528 void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
530 __rb_insert(node, root, augment_rotate);
532 EXPORT_SYMBOL(__rb_insert_augmented);
534 void rb_erase_augmented(struct rb_node *node, struct rb_root *root,
535 const struct rb_augment_callbacks *augment)
537 __rb_erase(node, root, augment);
539 EXPORT_SYMBOL(rb_erase_augmented);
541 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
543 struct rb_node *parent;
547 parent = rb_parent(node);
551 if (node == parent->rb_left && parent->rb_right)
552 func(parent->rb_right, data);
553 else if (parent->rb_left)
554 func(parent->rb_left, data);
561 * after inserting @node into the tree, update the tree to account for
562 * both the new entry and any damage done by rebalance
564 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
567 node = node->rb_left;
568 else if (node->rb_right)
569 node = node->rb_right;
571 rb_augment_path(node, func, data);
573 EXPORT_SYMBOL(rb_augment_insert);
576 * before removing the node, find the deepest node on the rebalance path
577 * that will still be there after @node gets removed
579 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
581 struct rb_node *deepest;
583 if (!node->rb_right && !node->rb_left)
584 deepest = rb_parent(node);
585 else if (!node->rb_right)
586 deepest = node->rb_left;
587 else if (!node->rb_left)
588 deepest = node->rb_right;
590 deepest = rb_next(node);
591 if (deepest->rb_right)
592 deepest = deepest->rb_right;
593 else if (rb_parent(deepest) != node)
594 deepest = rb_parent(deepest);
599 EXPORT_SYMBOL(rb_augment_erase_begin);
602 * after removal, update the tree to account for the removed entry
603 * and any rebalance damage.
605 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
608 rb_augment_path(node, func, data);
610 EXPORT_SYMBOL(rb_augment_erase_end);
613 * This function returns the first node (in sort order) of the tree.
615 struct rb_node *rb_first(const struct rb_root *root)
626 EXPORT_SYMBOL(rb_first);
628 struct rb_node *rb_last(const struct rb_root *root)
639 EXPORT_SYMBOL(rb_last);
641 struct rb_node *rb_next(const struct rb_node *node)
643 struct rb_node *parent;
645 if (RB_EMPTY_NODE(node))
649 * If we have a right-hand child, go down and then left as far
652 if (node->rb_right) {
653 node = node->rb_right;
654 while (node->rb_left)
656 return (struct rb_node *)node;
660 * No right-hand children. Everything down and left is smaller than us,
661 * so any 'next' node must be in the general direction of our parent.
662 * Go up the tree; any time the ancestor is a right-hand child of its
663 * parent, keep going up. First time it's a left-hand child of its
664 * parent, said parent is our 'next' node.
666 while ((parent = rb_parent(node)) && node == parent->rb_right)
671 EXPORT_SYMBOL(rb_next);
673 struct rb_node *rb_prev(const struct rb_node *node)
675 struct rb_node *parent;
677 if (RB_EMPTY_NODE(node))
681 * If we have a left-hand child, go down and then right as far
685 node = node->rb_left;
686 while (node->rb_right)
688 return (struct rb_node *)node;
692 * No left-hand children. Go up till we find an ancestor which
693 * is a right-hand child of its parent.
695 while ((parent = rb_parent(node)) && node == parent->rb_left)
700 EXPORT_SYMBOL(rb_prev);
702 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
703 struct rb_root *root)
705 struct rb_node *parent = rb_parent(victim);
707 /* Set the surrounding nodes to point to the replacement */
708 __rb_change_child(victim, new, parent, root);
710 rb_set_parent(victim->rb_left, new);
711 if (victim->rb_right)
712 rb_set_parent(victim->rb_right, new);
714 /* Copy the pointers/colour from the victim to the replacement */
717 EXPORT_SYMBOL(rb_replace_node);