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
28 #include "wimlib/rbtree.h"
32 * red-black trees properties: http://en.wikipedia.org/wiki/Rbtree
34 * 1) A node is either red or black
35 * 2) The root is black
36 * 3) All leaves (NULL) are black
37 * 4) Both children of every red node are black
38 * 5) Every simple path from root to leaves contains the same number
41 * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
42 * consecutive red nodes in a path and every red node is therefore followed by
43 * a black. So if B is the number of black nodes on every simple path (as per
44 * 5), then the longest possible path due to 4 is 2B.
46 * We shall indicate color with case, where black nodes are uppercase and red
47 * nodes will be lowercase. Unknown color nodes shall be drawn as red within
48 * parentheses and have some accompanying text comment.
51 struct rb_augment_callbacks {
52 void (*propagate)(struct rb_node *node, struct rb_node *stop);
53 void (*copy)(struct rb_node *old, struct rb_node *new);
54 void (*rotate)(struct rb_node *old, struct rb_node *new);
60 #define __rb_parent(pc) ((struct rb_node *)(pc & ~3))
62 #define __rb_color(pc) ((pc) & 1)
63 #define __rb_is_black(pc) __rb_color(pc)
64 #define __rb_is_red(pc) (!__rb_color(pc))
65 #define rb_color(rb) __rb_color((rb)->__rb_parent_color)
66 #define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
67 #define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)
69 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
71 rb->__rb_parent_color = rb_color(rb) | (uintptr_t)p;
74 static inline void rb_set_parent_color(struct rb_node *rb,
75 struct rb_node *p, int color)
77 rb->__rb_parent_color = (uintptr_t)p | color;
80 static inline void rb_set_black(struct rb_node *rb)
82 rb->__rb_parent_color |= RB_BLACK;
85 static inline struct rb_node *rb_red_parent(struct rb_node *red)
87 return (struct rb_node *)red->__rb_parent_color;
91 __rb_change_child(struct rb_node *old, struct rb_node *new,
92 struct rb_node *parent, struct rb_root *root)
95 if (parent->rb_left == old)
96 parent->rb_left = new;
98 parent->rb_right = new;
104 * Helper function for rotations:
105 * - old's parent and color get assigned to new
106 * - old gets assigned new as a parent and 'color' as a color.
109 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
110 struct rb_root *root, int color)
112 struct rb_node *parent = rb_parent(old);
113 new->__rb_parent_color = old->__rb_parent_color;
114 rb_set_parent_color(old, new, color);
115 __rb_change_child(old, new, parent, root);
119 __rb_erase_color(struct rb_node *parent, struct rb_root *root,
120 void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
122 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
127 * - node is black (or NULL on first iteration)
128 * - node is not the root (parent is not NULL)
129 * - All leaf paths going through parent and node have a
130 * black node count that is 1 lower than other leaf paths.
132 sibling = parent->rb_right;
133 if (node != sibling) { /* node == parent->rb_left */
134 if (rb_is_red(sibling)) {
136 * Case 1 - left rotate at parent
144 parent->rb_right = tmp1 = sibling->rb_left;
145 sibling->rb_left = parent;
146 rb_set_parent_color(tmp1, parent, RB_BLACK);
147 __rb_rotate_set_parents(parent, sibling, root,
149 augment_rotate(parent, sibling);
152 tmp1 = sibling->rb_right;
153 if (!tmp1 || rb_is_black(tmp1)) {
154 tmp2 = sibling->rb_left;
155 if (!tmp2 || rb_is_black(tmp2)) {
157 * Case 2 - sibling color flip
158 * (p could be either color here)
166 * This leaves us violating 5) which
167 * can be fixed by flipping p to black
168 * if it was red, or by recursing at p.
169 * p is red when coming from Case 1.
171 rb_set_parent_color(sibling, parent,
173 if (rb_is_red(parent))
174 rb_set_black(parent);
177 parent = rb_parent(node);
184 * Case 3 - right rotate at sibling
185 * (p could be either color here)
195 sibling->rb_left = tmp1 = tmp2->rb_right;
196 tmp2->rb_right = sibling;
197 parent->rb_right = tmp2;
199 rb_set_parent_color(tmp1, sibling,
201 augment_rotate(sibling, tmp2);
206 * Case 4 - left rotate at parent + color flips
207 * (p and sl could be either color here.
208 * After rotation, p becomes black, s acquires
209 * p's color, and sl keeps its color)
217 parent->rb_right = tmp2 = sibling->rb_left;
218 sibling->rb_left = parent;
219 rb_set_parent_color(tmp1, sibling, RB_BLACK);
221 rb_set_parent(tmp2, parent);
222 __rb_rotate_set_parents(parent, sibling, root,
224 augment_rotate(parent, sibling);
227 sibling = parent->rb_left;
228 if (rb_is_red(sibling)) {
229 /* Case 1 - right rotate at parent */
230 parent->rb_left = tmp1 = sibling->rb_right;
231 sibling->rb_right = parent;
232 rb_set_parent_color(tmp1, parent, RB_BLACK);
233 __rb_rotate_set_parents(parent, sibling, root,
235 augment_rotate(parent, sibling);
238 tmp1 = sibling->rb_left;
239 if (!tmp1 || rb_is_black(tmp1)) {
240 tmp2 = sibling->rb_right;
241 if (!tmp2 || rb_is_black(tmp2)) {
242 /* Case 2 - sibling color flip */
243 rb_set_parent_color(sibling, parent,
245 if (rb_is_red(parent))
246 rb_set_black(parent);
249 parent = rb_parent(node);
255 /* Case 3 - right rotate at sibling */
256 sibling->rb_right = tmp1 = tmp2->rb_left;
257 tmp2->rb_left = sibling;
258 parent->rb_left = tmp2;
260 rb_set_parent_color(tmp1, sibling,
262 augment_rotate(sibling, tmp2);
266 /* Case 4 - left rotate at parent + color flips */
267 parent->rb_left = tmp2 = sibling->rb_right;
268 sibling->rb_right = parent;
269 rb_set_parent_color(tmp1, sibling, RB_BLACK);
271 rb_set_parent(tmp2, parent);
272 __rb_rotate_set_parents(parent, sibling, root,
274 augment_rotate(parent, sibling);
281 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
282 const struct rb_augment_callbacks *augment)
284 struct rb_node *child = node->rb_right, *tmp = node->rb_left;
285 struct rb_node *parent, *rebalance;
290 * Case 1: node to erase has no more than 1 child (easy!)
292 * Note that if there is one child it must be red due to 5)
293 * and node must be black due to 4). We adjust colors locally
294 * so as to bypass __rb_erase_color() later on.
296 pc = node->__rb_parent_color;
297 parent = __rb_parent(pc);
298 __rb_change_child(node, child, parent, root);
300 child->__rb_parent_color = pc;
303 rebalance = __rb_is_black(pc) ? parent : NULL;
306 /* Still case 1, but this time the child is node->rb_left */
307 tmp->__rb_parent_color = pc = node->__rb_parent_color;
308 parent = __rb_parent(pc);
309 __rb_change_child(node, tmp, parent, root);
313 struct rb_node *successor = child, *child2;
314 tmp = child->rb_left;
317 * Case 2: node's successor is its right child
326 child2 = successor->rb_right;
327 augment->copy(node, successor);
330 * Case 3: node's successor is leftmost under
331 * node's right child subtree
348 parent->rb_left = child2 = successor->rb_right;
349 successor->rb_right = child;
350 rb_set_parent(child, successor);
351 augment->copy(node, successor);
352 augment->propagate(parent, successor);
355 successor->rb_left = tmp = node->rb_left;
356 rb_set_parent(tmp, successor);
358 pc = node->__rb_parent_color;
359 tmp = __rb_parent(pc);
360 __rb_change_child(node, successor, tmp, root);
362 successor->__rb_parent_color = pc;
363 rb_set_parent_color(child2, parent, RB_BLACK);
366 uintptr_t pc2 = successor->__rb_parent_color;
367 successor->__rb_parent_color = pc;
368 rebalance = __rb_is_black(pc2) ? parent : NULL;
373 augment->propagate(tmp, NULL);
375 __rb_erase_color(rebalance, root, augment->rotate);
380 __rb_insert(struct rb_node *node, struct rb_root *root,
381 void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
383 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
387 * Loop invariant: node is red
389 * If there is a black parent, we are done.
390 * Otherwise, take some corrective action as we don't
391 * want a red root or two consecutive red nodes.
394 rb_set_parent_color(node, NULL, RB_BLACK);
396 } else if (rb_is_black(parent))
399 gparent = rb_red_parent(parent);
401 tmp = gparent->rb_right;
402 if (parent != tmp) { /* parent == gparent->rb_left */
403 if (tmp && rb_is_red(tmp)) {
405 * Case 1 - color flips
413 * However, since g's parent might be red, and
414 * 4) does not allow this, we need to recurse
417 rb_set_parent_color(tmp, gparent, RB_BLACK);
418 rb_set_parent_color(parent, gparent, RB_BLACK);
420 parent = rb_parent(node);
421 rb_set_parent_color(node, parent, RB_RED);
425 tmp = parent->rb_right;
428 * Case 2 - left rotate at parent
436 * This still leaves us in violation of 4), the
437 * continuation into Case 3 will fix that.
439 parent->rb_right = tmp = node->rb_left;
440 node->rb_left = parent;
442 rb_set_parent_color(tmp, parent,
444 rb_set_parent_color(parent, node, RB_RED);
445 augment_rotate(parent, node);
447 tmp = node->rb_right;
451 * Case 3 - right rotate at gparent
459 gparent->rb_left = tmp; /* == parent->rb_right */
460 parent->rb_right = gparent;
462 rb_set_parent_color(tmp, gparent, RB_BLACK);
463 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
464 augment_rotate(gparent, parent);
467 tmp = gparent->rb_left;
468 if (tmp && rb_is_red(tmp)) {
469 /* Case 1 - color flips */
470 rb_set_parent_color(tmp, gparent, RB_BLACK);
471 rb_set_parent_color(parent, gparent, RB_BLACK);
473 parent = rb_parent(node);
474 rb_set_parent_color(node, parent, RB_RED);
478 tmp = parent->rb_left;
480 /* Case 2 - right rotate at parent */
481 parent->rb_left = tmp = node->rb_right;
482 node->rb_right = parent;
484 rb_set_parent_color(tmp, parent,
486 rb_set_parent_color(parent, node, RB_RED);
487 augment_rotate(parent, node);
492 /* Case 3 - left rotate at gparent */
493 gparent->rb_right = tmp; /* == parent->rb_left */
494 parent->rb_left = gparent;
496 rb_set_parent_color(tmp, gparent, RB_BLACK);
497 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
498 augment_rotate(gparent, parent);
506 * Non-augmented rbtree manipulation functions.
508 * We use dummy augmented callbacks here, and have the compiler optimize them
509 * out of the rb_insert_color() and rb_erase() function definitions.
512 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
513 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
514 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
516 static const struct rb_augment_callbacks dummy_callbacks = {
517 dummy_propagate, dummy_copy, dummy_rotate
520 void rb_insert_color(struct rb_node *node, struct rb_root *root)
522 __rb_insert(node, root, dummy_rotate);
525 void rb_erase(struct rb_node *node, struct rb_root *root)
527 rb_erase_augmented(node, root, &dummy_callbacks);