3 * Copyright (C) 1999 Andrea Arcangeli <andrea@suse.de>
4 * Copyright (C) 2002 David Woodhouse <dwmw2@infradead.org>
5 * Copyright (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
26 #include "wimlib/rbtree.h"
30 * red-black trees properties: http://en.wikipedia.org/wiki/Rbtree
32 * 1) A node is either red or black
33 * 2) The root is black
34 * 3) All leaves (NULL) are black
35 * 4) Both children of every red node are black
36 * 5) Every simple path from root to leaves contains the same number
39 * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
40 * consecutive red nodes in a path and every red node is therefore followed by
41 * a black. So if B is the number of black nodes on every simple path (as per
42 * 5), then the longest possible path due to 4 is 2B.
44 * We shall indicate color with case, where black nodes are uppercase and red
45 * nodes will be lowercase. Unknown color nodes shall be drawn as red within
46 * parentheses and have some accompanying text comment.
52 #define __rb_parent(pc) ((struct rb_node *)(pc & ~1))
54 #define __rb_color(pc) ((pc) & 1)
55 #define __rb_is_black(pc) __rb_color(pc)
56 #define __rb_is_red(pc) (!__rb_color(pc))
57 #define rb_color(rb) __rb_color((rb)->__rb_parent_color)
58 #define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
59 #define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)
62 rb_set_parent(struct rb_node *rb, struct rb_node *p)
64 rb->__rb_parent_color = rb_color(rb) | (uintptr_t)p;
68 rb_set_parent_color(struct rb_node *rb, struct rb_node *p, int color)
70 rb->__rb_parent_color = (uintptr_t)p | color;
74 rb_set_black(struct rb_node *rb)
76 rb->__rb_parent_color |= RB_BLACK;
79 static inline struct rb_node *
80 rb_red_parent(struct rb_node *red)
82 return (struct rb_node *)red->__rb_parent_color;
86 rb_change_child(struct rb_node *old, struct rb_node *new,
87 struct rb_node *parent, struct rb_root *root)
90 if (parent->rb_left == old)
91 parent->rb_left = new;
93 parent->rb_right = new;
99 * Helper function for rotations:
100 * - old's parent and color get assigned to new
101 * - old gets assigned new as a parent and 'color' as a color.
104 rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
105 struct rb_root *root, int color)
107 struct rb_node *parent = rb_parent(old);
108 new->__rb_parent_color = old->__rb_parent_color;
109 rb_set_parent_color(old, new, color);
110 rb_change_child(old, new, parent, root);
114 rb_erase_color(struct rb_node *parent, struct rb_root *root)
116 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
121 * - node is black (or NULL on first iteration)
122 * - node is not the root (parent is not NULL)
123 * - All leaf paths going through parent and node have a
124 * black node count that is 1 lower than other leaf paths.
126 sibling = parent->rb_right;
127 if (node != sibling) { /* node == parent->rb_left */
128 if (rb_is_red(sibling)) {
130 * Case 1 - left rotate at parent
138 parent->rb_right = tmp1 = sibling->rb_left;
139 sibling->rb_left = parent;
140 rb_set_parent_color(tmp1, parent, RB_BLACK);
141 rb_rotate_set_parents(parent, sibling, root, RB_RED);
144 tmp1 = sibling->rb_right;
145 if (!tmp1 || rb_is_black(tmp1)) {
146 tmp2 = sibling->rb_left;
147 if (!tmp2 || rb_is_black(tmp2)) {
149 * Case 2 - sibling color flip
150 * (p could be either color here)
158 * This leaves us violating 5) which
159 * can be fixed by flipping p to black
160 * if it was red, or by recursing at p.
161 * p is red when coming from Case 1.
163 rb_set_parent_color(sibling, parent,
165 if (rb_is_red(parent))
166 rb_set_black(parent);
169 parent = rb_parent(node);
176 * Case 3 - right rotate at sibling
177 * (p could be either color here)
187 sibling->rb_left = tmp1 = tmp2->rb_right;
188 tmp2->rb_right = sibling;
189 parent->rb_right = tmp2;
191 rb_set_parent_color(tmp1, sibling,
197 * Case 4 - left rotate at parent + color flips
198 * (p and sl could be either color here.
199 * After rotation, p becomes black, s acquires
200 * p's color, and sl keeps its color)
208 parent->rb_right = tmp2 = sibling->rb_left;
209 sibling->rb_left = parent;
210 rb_set_parent_color(tmp1, sibling, RB_BLACK);
212 rb_set_parent(tmp2, parent);
213 rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
216 sibling = parent->rb_left;
217 if (rb_is_red(sibling)) {
218 /* Case 1 - right rotate at parent */
219 parent->rb_left = tmp1 = sibling->rb_right;
220 sibling->rb_right = parent;
221 rb_set_parent_color(tmp1, parent, RB_BLACK);
222 rb_rotate_set_parents(parent, sibling, root,
226 tmp1 = sibling->rb_left;
227 if (!tmp1 || rb_is_black(tmp1)) {
228 tmp2 = sibling->rb_right;
229 if (!tmp2 || rb_is_black(tmp2)) {
230 /* Case 2 - sibling color flip */
231 rb_set_parent_color(sibling, parent,
233 if (rb_is_red(parent))
234 rb_set_black(parent);
237 parent = rb_parent(node);
243 /* Case 3 - right rotate at sibling */
244 sibling->rb_right = tmp1 = tmp2->rb_left;
245 tmp2->rb_left = sibling;
246 parent->rb_left = tmp2;
248 rb_set_parent_color(tmp1, sibling,
253 /* Case 4 - left rotate at parent + color flips */
254 parent->rb_left = tmp2 = sibling->rb_right;
255 sibling->rb_right = parent;
256 rb_set_parent_color(tmp1, sibling, RB_BLACK);
258 rb_set_parent(tmp2, parent);
259 rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
266 rb_erase(struct rb_node *node, struct rb_root *root)
268 struct rb_node *child = node->rb_right, *tmp = node->rb_left;
269 struct rb_node *parent, *rebalance;
274 * Case 1: node to erase has no more than 1 child (easy!)
276 * Note that if there is one child it must be red due to 5)
277 * and node must be black due to 4). We adjust colors locally
278 * so as to bypass __rb_erase_color() later on.
280 pc = node->__rb_parent_color;
281 parent = __rb_parent(pc);
282 rb_change_child(node, child, parent, root);
284 child->__rb_parent_color = pc;
287 rebalance = __rb_is_black(pc) ? parent : NULL;
290 /* Still case 1, but this time the child is node->rb_left */
291 tmp->__rb_parent_color = pc = node->__rb_parent_color;
292 parent = __rb_parent(pc);
293 rb_change_child(node, tmp, parent, root);
297 struct rb_node *successor = child, *child2;
298 tmp = child->rb_left;
301 * Case 2: node's successor is its right child
310 child2 = successor->rb_right;
313 * Case 3: node's successor is leftmost under
314 * node's right child subtree
331 parent->rb_left = child2 = successor->rb_right;
332 successor->rb_right = child;
333 rb_set_parent(child, successor);
336 successor->rb_left = tmp = node->rb_left;
337 rb_set_parent(tmp, successor);
339 pc = node->__rb_parent_color;
340 tmp = __rb_parent(pc);
341 rb_change_child(node, successor, tmp, root);
343 successor->__rb_parent_color = pc;
344 rb_set_parent_color(child2, parent, RB_BLACK);
347 uintptr_t pc2 = successor->__rb_parent_color;
348 successor->__rb_parent_color = pc;
349 rebalance = __rb_is_black(pc2) ? parent : NULL;
355 rb_erase_color(rebalance, root);
359 rb_insert_color(struct rb_node *node, struct rb_root *root)
361 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
365 * Loop invariant: node is red
367 * If there is a black parent, we are done.
368 * Otherwise, take some corrective action as we don't
369 * want a red root or two consecutive red nodes.
372 rb_set_parent_color(node, NULL, RB_BLACK);
374 } else if (rb_is_black(parent))
377 gparent = rb_red_parent(parent);
379 tmp = gparent->rb_right;
380 if (parent != tmp) { /* parent == gparent->rb_left */
381 if (tmp && rb_is_red(tmp)) {
383 * Case 1 - color flips
391 * However, since g's parent might be red, and
392 * 4) does not allow this, we need to recurse
395 rb_set_parent_color(tmp, gparent, RB_BLACK);
396 rb_set_parent_color(parent, gparent, RB_BLACK);
398 parent = rb_parent(node);
399 rb_set_parent_color(node, parent, RB_RED);
403 tmp = parent->rb_right;
406 * Case 2 - left rotate at parent
414 * This still leaves us in violation of 4), the
415 * continuation into Case 3 will fix that.
417 parent->rb_right = tmp = node->rb_left;
418 node->rb_left = parent;
420 rb_set_parent_color(tmp, parent,
422 rb_set_parent_color(parent, node, RB_RED);
424 tmp = node->rb_right;
428 * Case 3 - right rotate at gparent
436 gparent->rb_left = tmp; /* == parent->rb_right */
437 parent->rb_right = gparent;
439 rb_set_parent_color(tmp, gparent, RB_BLACK);
440 rb_rotate_set_parents(gparent, parent, root, RB_RED);
443 tmp = gparent->rb_left;
444 if (tmp && rb_is_red(tmp)) {
445 /* Case 1 - color flips */
446 rb_set_parent_color(tmp, gparent, RB_BLACK);
447 rb_set_parent_color(parent, gparent, RB_BLACK);
449 parent = rb_parent(node);
450 rb_set_parent_color(node, parent, RB_RED);
454 tmp = parent->rb_left;
456 /* Case 2 - right rotate at parent */
457 parent->rb_left = tmp = node->rb_right;
458 node->rb_right = parent;
460 rb_set_parent_color(tmp, parent,
462 rb_set_parent_color(parent, node, RB_RED);
467 /* Case 3 - left rotate at gparent */
468 gparent->rb_right = tmp; /* == parent->rb_left */
469 parent->rb_left = gparent;
471 rb_set_parent_color(tmp, gparent, RB_BLACK);
472 rb_rotate_set_parents(gparent, parent, root, RB_RED);
478 static struct rb_node *
479 rb_left_deepest_node(const struct rb_node *node)
483 node = node->rb_left;
484 else if (node->rb_right)
485 node = node->rb_right;
487 return (struct rb_node *)node;
492 rb_next_postorder(const struct rb_node *node, const struct rb_node *parent)
494 /* If we're sitting on node, we've already seen our children */
495 if (parent && node == parent->rb_left && parent->rb_right) {
496 /* If we are the parent's left node, go to the parent's right
497 * node then all the way down to the left */
498 return rb_left_deepest_node(parent->rb_right);
500 /* Otherwise we are the parent's right node, and the parent
502 return (struct rb_node *)parent;
506 rb_first_postorder(const struct rb_root *root)
511 return rb_left_deepest_node(root->rb_node);
515 rb_next(const struct rb_node *node)
517 struct rb_node *parent;
520 * If we have a right-hand child, go down and then left as far
523 if (node->rb_right) {
524 node = node->rb_right;
525 while (node->rb_left)
526 node = node->rb_left;
527 return (struct rb_node *)node;
531 * No right-hand children. Everything down and left is smaller than us,
532 * so any 'next' node must be in the general direction of our parent.
533 * Go up the tree; any time the ancestor is a right-hand child of its
534 * parent, keep going up. First time it's a left-hand child of its
535 * parent, said parent is our 'next' node.
537 while ((parent = rb_parent(node)) && node == parent->rb_right)
544 * This function returns the first node (in sort order) of the tree.
547 rb_first(const struct rb_root *root)