Refactor headers
[wimlib] / src / rbtree.c
1 /*
2   Red Black Trees
3   (C) 1999  Andrea Arcangeli <andrea@suse.de>
4   (C) 2002  David Woodhouse <dwmw2@infradead.org>
5   (C) 2012  Michel Lespinasse <walken@google.com>
6
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.
11
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.
16
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
20
21   linux/lib/rbtree.c
22 */
23
24 #ifdef HAVE_CONFIG_H
25 #  include "config.h"
26 #endif
27
28 #include "wimlib/rbtree.h"
29 #include <stdbool.h>
30
31 /*
32  * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
33  *
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
39  *     of black nodes.
40  *
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.
45  *
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.
49  */
50
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);
55 };
56
57 #define RB_RED          0
58 #define RB_BLACK        1
59
60 #define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))
61
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)
68
69 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
70 {
71         rb->__rb_parent_color = rb_color(rb) | (uintptr_t)p;
72 }
73
74 static inline void rb_set_parent_color(struct rb_node *rb,
75                                        struct rb_node *p, int color)
76 {
77         rb->__rb_parent_color = (uintptr_t)p | color;
78 }
79
80 static inline void rb_set_black(struct rb_node *rb)
81 {
82         rb->__rb_parent_color |= RB_BLACK;
83 }
84
85 static inline struct rb_node *rb_red_parent(struct rb_node *red)
86 {
87         return (struct rb_node *)red->__rb_parent_color;
88 }
89
90 static inline void
91 __rb_change_child(struct rb_node *old, struct rb_node *new,
92                   struct rb_node *parent, struct rb_root *root)
93 {
94         if (parent) {
95                 if (parent->rb_left == old)
96                         parent->rb_left = new;
97                 else
98                         parent->rb_right = new;
99         } else
100                 root->rb_node = new;
101 }
102
103 /*
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.
107  */
108 static inline void
109 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
110                         struct rb_root *root, int color)
111 {
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);
116 }
117
118 static inline void
119 __rb_erase_color(struct rb_node *parent, struct rb_root *root,
120         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
121 {
122         struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
123
124         while (true) {
125                 /*
126                  * Loop invariants:
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.
131                  */
132                 sibling = parent->rb_right;
133                 if (node != sibling) {  /* node == parent->rb_left */
134                         if (rb_is_red(sibling)) {
135                                 /*
136                                  * Case 1 - left rotate at parent
137                                  *
138                                  *     P               S
139                                  *    / \             / \
140                                  *   N   s    -->    p   Sr
141                                  *      / \         / \
142                                  *     Sl  Sr      N   Sl
143                                  */
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,
148                                                         RB_RED);
149                                 augment_rotate(parent, sibling);
150                                 sibling = tmp1;
151                         }
152                         tmp1 = sibling->rb_right;
153                         if (!tmp1 || rb_is_black(tmp1)) {
154                                 tmp2 = sibling->rb_left;
155                                 if (!tmp2 || rb_is_black(tmp2)) {
156                                         /*
157                                          * Case 2 - sibling color flip
158                                          * (p could be either color here)
159                                          *
160                                          *    (p)           (p)
161                                          *    / \           / \
162                                          *   N   S    -->  N   s
163                                          *      / \           / \
164                                          *     Sl  Sr        Sl  Sr
165                                          *
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.
170                                          */
171                                         rb_set_parent_color(sibling, parent,
172                                                             RB_RED);
173                                         if (rb_is_red(parent))
174                                                 rb_set_black(parent);
175                                         else {
176                                                 node = parent;
177                                                 parent = rb_parent(node);
178                                                 if (parent)
179                                                         continue;
180                                         }
181                                         break;
182                                 }
183                                 /*
184                                  * Case 3 - right rotate at sibling
185                                  * (p could be either color here)
186                                  *
187                                  *   (p)           (p)
188                                  *   / \           / \
189                                  *  N   S    -->  N   Sl
190                                  *     / \             \
191                                  *    sl  Sr            s
192                                  *                       \
193                                  *                        Sr
194                                  */
195                                 sibling->rb_left = tmp1 = tmp2->rb_right;
196                                 tmp2->rb_right = sibling;
197                                 parent->rb_right = tmp2;
198                                 if (tmp1)
199                                         rb_set_parent_color(tmp1, sibling,
200                                                             RB_BLACK);
201                                 augment_rotate(sibling, tmp2);
202                                 tmp1 = sibling;
203                                 sibling = tmp2;
204                         }
205                         /*
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)
210                          *
211                          *      (p)             (s)
212                          *      / \             / \
213                          *     N   S     -->   P   Sr
214                          *        / \         / \
215                          *      (sl) sr      N  (sl)
216                          */
217                         parent->rb_right = tmp2 = sibling->rb_left;
218                         sibling->rb_left = parent;
219                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
220                         if (tmp2)
221                                 rb_set_parent(tmp2, parent);
222                         __rb_rotate_set_parents(parent, sibling, root,
223                                                 RB_BLACK);
224                         augment_rotate(parent, sibling);
225                         break;
226                 } else {
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,
234                                                         RB_RED);
235                                 augment_rotate(parent, sibling);
236                                 sibling = tmp1;
237                         }
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,
244                                                             RB_RED);
245                                         if (rb_is_red(parent))
246                                                 rb_set_black(parent);
247                                         else {
248                                                 node = parent;
249                                                 parent = rb_parent(node);
250                                                 if (parent)
251                                                         continue;
252                                         }
253                                         break;
254                                 }
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;
259                                 if (tmp1)
260                                         rb_set_parent_color(tmp1, sibling,
261                                                             RB_BLACK);
262                                 augment_rotate(sibling, tmp2);
263                                 tmp1 = sibling;
264                                 sibling = tmp2;
265                         }
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);
270                         if (tmp2)
271                                 rb_set_parent(tmp2, parent);
272                         __rb_rotate_set_parents(parent, sibling, root,
273                                                 RB_BLACK);
274                         augment_rotate(parent, sibling);
275                         break;
276                 }
277         }
278 }
279
280 static inline void
281 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
282                    const struct rb_augment_callbacks *augment)
283 {
284         struct rb_node *child = node->rb_right, *tmp = node->rb_left;
285         struct rb_node *parent, *rebalance;
286         uintptr_t pc;
287
288         if (!tmp) {
289                 /*
290                  * Case 1: node to erase has no more than 1 child (easy!)
291                  *
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.
295                  */
296                 pc = node->__rb_parent_color;
297                 parent = __rb_parent(pc);
298                 __rb_change_child(node, child, parent, root);
299                 if (child) {
300                         child->__rb_parent_color = pc;
301                         rebalance = NULL;
302                 } else
303                         rebalance = __rb_is_black(pc) ? parent : NULL;
304                 tmp = parent;
305         } else if (!child) {
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);
310                 rebalance = NULL;
311                 tmp = parent;
312         } else {
313                 struct rb_node *successor = child, *child2;
314                 tmp = child->rb_left;
315                 if (!tmp) {
316                         /*
317                          * Case 2: node's successor is its right child
318                          *
319                          *    (n)          (s)
320                          *    / \          / \
321                          *  (x) (s)  ->  (x) (c)
322                          *        \
323                          *        (c)
324                          */
325                         parent = successor;
326                         child2 = successor->rb_right;
327                         augment->copy(node, successor);
328                 } else {
329                         /*
330                          * Case 3: node's successor is leftmost under
331                          * node's right child subtree
332                          *
333                          *    (n)          (s)
334                          *    / \          / \
335                          *  (x) (y)  ->  (x) (y)
336                          *      /            /
337                          *    (p)          (p)
338                          *    /            /
339                          *  (s)          (c)
340                          *    \
341                          *    (c)
342                          */
343                         do {
344                                 parent = successor;
345                                 successor = tmp;
346                                 tmp = tmp->rb_left;
347                         } while (tmp);
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);
353                 }
354
355                 successor->rb_left = tmp = node->rb_left;
356                 rb_set_parent(tmp, successor);
357
358                 pc = node->__rb_parent_color;
359                 tmp = __rb_parent(pc);
360                 __rb_change_child(node, successor, tmp, root);
361                 if (child2) {
362                         successor->__rb_parent_color = pc;
363                         rb_set_parent_color(child2, parent, RB_BLACK);
364                         rebalance = NULL;
365                 } else {
366                         uintptr_t pc2 = successor->__rb_parent_color;
367                         successor->__rb_parent_color = pc;
368                         rebalance = __rb_is_black(pc2) ? parent : NULL;
369                 }
370                 tmp = successor;
371         }
372
373         augment->propagate(tmp, NULL);
374         if (rebalance)
375                 __rb_erase_color(rebalance, root, augment->rotate);
376 }
377
378
379 static inline void
380 __rb_insert(struct rb_node *node, struct rb_root *root,
381             void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
382 {
383         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
384
385         while (true) {
386                 /*
387                  * Loop invariant: node is red
388                  *
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.
392                  */
393                 if (!parent) {
394                         rb_set_parent_color(node, NULL, RB_BLACK);
395                         break;
396                 } else if (rb_is_black(parent))
397                         break;
398
399                 gparent = rb_red_parent(parent);
400
401                 tmp = gparent->rb_right;
402                 if (parent != tmp) {    /* parent == gparent->rb_left */
403                         if (tmp && rb_is_red(tmp)) {
404                                 /*
405                                  * Case 1 - color flips
406                                  *
407                                  *       G            g
408                                  *      / \          / \
409                                  *     p   u  -->   P   U
410                                  *    /            /
411                                  *   n            N
412                                  *
413                                  * However, since g's parent might be red, and
414                                  * 4) does not allow this, we need to recurse
415                                  * at g.
416                                  */
417                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
418                                 rb_set_parent_color(parent, gparent, RB_BLACK);
419                                 node = gparent;
420                                 parent = rb_parent(node);
421                                 rb_set_parent_color(node, parent, RB_RED);
422                                 continue;
423                         }
424
425                         tmp = parent->rb_right;
426                         if (node == tmp) {
427                                 /*
428                                  * Case 2 - left rotate at parent
429                                  *
430                                  *      G             G
431                                  *     / \           / \
432                                  *    p   U  -->    n   U
433                                  *     \           /
434                                  *      n         p
435                                  *
436                                  * This still leaves us in violation of 4), the
437                                  * continuation into Case 3 will fix that.
438                                  */
439                                 parent->rb_right = tmp = node->rb_left;
440                                 node->rb_left = parent;
441                                 if (tmp)
442                                         rb_set_parent_color(tmp, parent,
443                                                             RB_BLACK);
444                                 rb_set_parent_color(parent, node, RB_RED);
445                                 augment_rotate(parent, node);
446                                 parent = node;
447                                 tmp = node->rb_right;
448                         }
449
450                         /*
451                          * Case 3 - right rotate at gparent
452                          *
453                          *        G           P
454                          *       / \         / \
455                          *      p   U  -->  n   g
456                          *     /                 \
457                          *    n                   U
458                          */
459                         gparent->rb_left = tmp;  /* == parent->rb_right */
460                         parent->rb_right = gparent;
461                         if (tmp)
462                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
463                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
464                         augment_rotate(gparent, parent);
465                         break;
466                 } else {
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);
472                                 node = gparent;
473                                 parent = rb_parent(node);
474                                 rb_set_parent_color(node, parent, RB_RED);
475                                 continue;
476                         }
477
478                         tmp = parent->rb_left;
479                         if (node == tmp) {
480                                 /* Case 2 - right rotate at parent */
481                                 parent->rb_left = tmp = node->rb_right;
482                                 node->rb_right = parent;
483                                 if (tmp)
484                                         rb_set_parent_color(tmp, parent,
485                                                             RB_BLACK);
486                                 rb_set_parent_color(parent, node, RB_RED);
487                                 augment_rotate(parent, node);
488                                 parent = node;
489                                 tmp = node->rb_left;
490                         }
491
492                         /* Case 3 - left rotate at gparent */
493                         gparent->rb_right = tmp;  /* == parent->rb_left */
494                         parent->rb_left = gparent;
495                         if (tmp)
496                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
497                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
498                         augment_rotate(gparent, parent);
499                         break;
500                 }
501         }
502 }
503
504
505 /*
506  * Non-augmented rbtree manipulation functions.
507  *
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.
510  */
511
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) {}
515
516 static const struct rb_augment_callbacks dummy_callbacks = {
517         dummy_propagate, dummy_copy, dummy_rotate
518 };
519
520 void rb_insert_color(struct rb_node *node, struct rb_root *root)
521 {
522         __rb_insert(node, root, dummy_rotate);
523 }
524
525 void rb_erase(struct rb_node *node, struct rb_root *root)
526 {
527         rb_erase_augmented(node, root, &dummy_callbacks);
528 }