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 void rb_insert_color(struct rb_node *node, struct rb_root *root)
110 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
114 * Loop invariant: node is red
116 * If there is a black parent, we are done.
117 * Otherwise, take some corrective action as we don't
118 * want a red root or two consecutive red nodes.
121 rb_set_parent_color(node, NULL, RB_BLACK);
123 } else if (rb_is_black(parent))
126 gparent = rb_red_parent(parent);
128 tmp = gparent->rb_right;
129 if (parent != tmp) { /* parent == gparent->rb_left */
130 if (tmp && rb_is_red(tmp)) {
132 * Case 1 - color flips
140 * However, since g's parent might be red, and
141 * 4) does not allow this, we need to recurse
144 rb_set_parent_color(tmp, gparent, RB_BLACK);
145 rb_set_parent_color(parent, gparent, RB_BLACK);
147 parent = rb_parent(node);
148 rb_set_parent_color(node, parent, RB_RED);
152 tmp = parent->rb_right;
155 * Case 2 - left rotate at parent
163 * This still leaves us in violation of 4), the
164 * continuation into Case 3 will fix that.
166 parent->rb_right = tmp = node->rb_left;
167 node->rb_left = parent;
169 rb_set_parent_color(tmp, parent,
171 rb_set_parent_color(parent, node, RB_RED);
173 tmp = node->rb_right;
177 * Case 3 - right rotate at gparent
185 gparent->rb_left = tmp; /* == parent->rb_right */
186 parent->rb_right = gparent;
188 rb_set_parent_color(tmp, gparent, RB_BLACK);
189 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
192 tmp = gparent->rb_left;
193 if (tmp && rb_is_red(tmp)) {
194 /* Case 1 - color flips */
195 rb_set_parent_color(tmp, gparent, RB_BLACK);
196 rb_set_parent_color(parent, gparent, RB_BLACK);
198 parent = rb_parent(node);
199 rb_set_parent_color(node, parent, RB_RED);
203 tmp = parent->rb_left;
205 /* Case 2 - right rotate at parent */
206 parent->rb_left = tmp = node->rb_right;
207 node->rb_right = parent;
209 rb_set_parent_color(tmp, parent,
211 rb_set_parent_color(parent, node, RB_RED);
216 /* Case 3 - left rotate at gparent */
217 gparent->rb_right = tmp; /* == parent->rb_left */
218 parent->rb_left = gparent;
220 rb_set_parent_color(tmp, gparent, RB_BLACK);
221 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
226 EXPORT_SYMBOL(rb_insert_color);
228 static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
230 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
235 * - node is black (or NULL on first iteration)
236 * - node is not the root (parent is not NULL)
237 * - All leaf paths going through parent and node have a
238 * black node count that is 1 lower than other leaf paths.
240 sibling = parent->rb_right;
241 if (node != sibling) { /* node == parent->rb_left */
242 if (rb_is_red(sibling)) {
244 * Case 1 - left rotate at parent
252 parent->rb_right = tmp1 = sibling->rb_left;
253 sibling->rb_left = parent;
254 rb_set_parent_color(tmp1, parent, RB_BLACK);
255 __rb_rotate_set_parents(parent, sibling, root,
259 tmp1 = sibling->rb_right;
260 if (!tmp1 || rb_is_black(tmp1)) {
261 tmp2 = sibling->rb_left;
262 if (!tmp2 || rb_is_black(tmp2)) {
264 * Case 2 - sibling color flip
265 * (p could be either color here)
273 * This leaves us violating 5) which
274 * can be fixed by flipping p to black
275 * if it was red, or by recursing at p.
276 * p is red when coming from Case 1.
278 rb_set_parent_color(sibling, parent,
280 if (rb_is_red(parent))
281 rb_set_black(parent);
284 parent = rb_parent(node);
291 * Case 3 - right rotate at sibling
292 * (p could be either color here)
302 sibling->rb_left = tmp1 = tmp2->rb_right;
303 tmp2->rb_right = sibling;
304 parent->rb_right = tmp2;
306 rb_set_parent_color(tmp1, sibling,
312 * Case 4 - left rotate at parent + color flips
313 * (p and sl could be either color here.
314 * After rotation, p becomes black, s acquires
315 * p's color, and sl keeps its color)
323 parent->rb_right = tmp2 = sibling->rb_left;
324 sibling->rb_left = parent;
325 rb_set_parent_color(tmp1, sibling, RB_BLACK);
327 rb_set_parent(tmp2, parent);
328 __rb_rotate_set_parents(parent, sibling, root,
332 sibling = parent->rb_left;
333 if (rb_is_red(sibling)) {
334 /* Case 1 - right rotate at parent */
335 parent->rb_left = tmp1 = sibling->rb_right;
336 sibling->rb_right = parent;
337 rb_set_parent_color(tmp1, parent, RB_BLACK);
338 __rb_rotate_set_parents(parent, sibling, root,
342 tmp1 = sibling->rb_left;
343 if (!tmp1 || rb_is_black(tmp1)) {
344 tmp2 = sibling->rb_right;
345 if (!tmp2 || rb_is_black(tmp2)) {
346 /* Case 2 - sibling color flip */
347 rb_set_parent_color(sibling, parent,
349 if (rb_is_red(parent))
350 rb_set_black(parent);
353 parent = rb_parent(node);
359 /* Case 3 - right rotate at sibling */
360 sibling->rb_right = tmp1 = tmp2->rb_left;
361 tmp2->rb_left = sibling;
362 parent->rb_left = tmp2;
364 rb_set_parent_color(tmp1, sibling,
369 /* Case 4 - left rotate at parent + color flips */
370 parent->rb_left = tmp2 = sibling->rb_right;
371 sibling->rb_right = parent;
372 rb_set_parent_color(tmp1, sibling, RB_BLACK);
374 rb_set_parent(tmp2, parent);
375 __rb_rotate_set_parents(parent, sibling, root,
382 void rb_erase(struct rb_node *node, struct rb_root *root)
384 struct rb_node *child = node->rb_right, *tmp = node->rb_left;
385 struct rb_node *parent, *rebalance;
390 * Case 1: node to erase has no more than 1 child (easy!)
392 * Note that if there is one child it must be red due to 5)
393 * and node must be black due to 4). We adjust colors locally
394 * so as to bypass __rb_erase_color() later on.
396 pc = node->__rb_parent_color;
397 parent = __rb_parent(pc);
398 __rb_change_child(node, child, parent, root);
400 child->__rb_parent_color = pc;
403 rebalance = __rb_is_black(pc) ? parent : NULL;
405 /* Still case 1, but this time the child is node->rb_left */
406 tmp->__rb_parent_color = pc = node->__rb_parent_color;
407 parent = __rb_parent(pc);
408 __rb_change_child(node, tmp, parent, root);
411 struct rb_node *successor = child, *child2;
412 tmp = child->rb_left;
415 * Case 2: node's successor is its right child
424 child2 = child->rb_right;
427 * Case 3: node's successor is leftmost under
428 * node's right child subtree
445 parent->rb_left = child2 = successor->rb_right;
446 successor->rb_right = child;
447 rb_set_parent(child, successor);
450 successor->rb_left = tmp = node->rb_left;
451 rb_set_parent(tmp, successor);
453 pc = node->__rb_parent_color;
454 tmp = __rb_parent(pc);
455 __rb_change_child(node, successor, tmp, root);
457 successor->__rb_parent_color = pc;
458 rb_set_parent_color(child2, parent, RB_BLACK);
461 unsigned long pc2 = successor->__rb_parent_color;
462 successor->__rb_parent_color = pc;
463 rebalance = __rb_is_black(pc2) ? parent : NULL;
468 __rb_erase_color(rebalance, root);
470 EXPORT_SYMBOL(rb_erase);
472 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
474 struct rb_node *parent;
478 parent = rb_parent(node);
482 if (node == parent->rb_left && parent->rb_right)
483 func(parent->rb_right, data);
484 else if (parent->rb_left)
485 func(parent->rb_left, data);
492 * after inserting @node into the tree, update the tree to account for
493 * both the new entry and any damage done by rebalance
495 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
498 node = node->rb_left;
499 else if (node->rb_right)
500 node = node->rb_right;
502 rb_augment_path(node, func, data);
504 EXPORT_SYMBOL(rb_augment_insert);
507 * before removing the node, find the deepest node on the rebalance path
508 * that will still be there after @node gets removed
510 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
512 struct rb_node *deepest;
514 if (!node->rb_right && !node->rb_left)
515 deepest = rb_parent(node);
516 else if (!node->rb_right)
517 deepest = node->rb_left;
518 else if (!node->rb_left)
519 deepest = node->rb_right;
521 deepest = rb_next(node);
522 if (deepest->rb_right)
523 deepest = deepest->rb_right;
524 else if (rb_parent(deepest) != node)
525 deepest = rb_parent(deepest);
530 EXPORT_SYMBOL(rb_augment_erase_begin);
533 * after removal, update the tree to account for the removed entry
534 * and any rebalance damage.
536 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
539 rb_augment_path(node, func, data);
541 EXPORT_SYMBOL(rb_augment_erase_end);
544 * This function returns the first node (in sort order) of the tree.
546 struct rb_node *rb_first(const struct rb_root *root)
557 EXPORT_SYMBOL(rb_first);
559 struct rb_node *rb_last(const struct rb_root *root)
570 EXPORT_SYMBOL(rb_last);
572 struct rb_node *rb_next(const struct rb_node *node)
574 struct rb_node *parent;
576 if (RB_EMPTY_NODE(node))
580 * If we have a right-hand child, go down and then left as far
583 if (node->rb_right) {
584 node = node->rb_right;
585 while (node->rb_left)
587 return (struct rb_node *)node;
591 * No right-hand children. Everything down and left is smaller than us,
592 * so any 'next' node must be in the general direction of our parent.
593 * Go up the tree; any time the ancestor is a right-hand child of its
594 * parent, keep going up. First time it's a left-hand child of its
595 * parent, said parent is our 'next' node.
597 while ((parent = rb_parent(node)) && node == parent->rb_right)
602 EXPORT_SYMBOL(rb_next);
604 struct rb_node *rb_prev(const struct rb_node *node)
606 struct rb_node *parent;
608 if (RB_EMPTY_NODE(node))
612 * If we have a left-hand child, go down and then right as far
616 node = node->rb_left;
617 while (node->rb_right)
619 return (struct rb_node *)node;
623 * No left-hand children. Go up till we find an ancestor which
624 * is a right-hand child of its parent.
626 while ((parent = rb_parent(node)) && node == parent->rb_left)
631 EXPORT_SYMBOL(rb_prev);
633 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
634 struct rb_root *root)
636 struct rb_node *parent = rb_parent(victim);
638 /* Set the surrounding nodes to point to the replacement */
639 __rb_change_child(victim, new, parent, root);
641 rb_set_parent(victim->rb_left, new);
642 if (victim->rb_right)
643 rb_set_parent(victim->rb_right, new);
645 /* Copy the pointers/colour from the victim to the replacement */
648 EXPORT_SYMBOL(rb_replace_node);