Root/lib/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#include <linux/rbtree_augmented.h>
25#include <linux/export.h>
26
27/*
28 * red-black trees properties: http://en.wikipedia.org/wiki/Rbtree
29 *
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
35 * of black nodes.
36 *
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.
41 *
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.
45 */
46
47static inline void rb_set_black(struct rb_node *rb)
48{
49    rb->__rb_parent_color |= RB_BLACK;
50}
51
52static inline struct rb_node *rb_red_parent(struct rb_node *red)
53{
54    return (struct rb_node *)red->__rb_parent_color;
55}
56
57/*
58 * Helper function for rotations:
59 * - old's parent and color get assigned to new
60 * - old gets assigned new as a parent and 'color' as a color.
61 */
62static inline void
63__rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
64            struct rb_root *root, int color)
65{
66    struct rb_node *parent = rb_parent(old);
67    new->__rb_parent_color = old->__rb_parent_color;
68    rb_set_parent_color(old, new, color);
69    __rb_change_child(old, new, parent, root);
70}
71
72static __always_inline void
73__rb_insert(struct rb_node *node, struct rb_root *root,
74        void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
75{
76    struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
77
78    while (true) {
79        /*
80         * Loop invariant: node is red
81         *
82         * If there is a black parent, we are done.
83         * Otherwise, take some corrective action as we don't
84         * want a red root or two consecutive red nodes.
85         */
86        if (!parent) {
87            rb_set_parent_color(node, NULL, RB_BLACK);
88            break;
89        } else if (rb_is_black(parent))
90            break;
91
92        gparent = rb_red_parent(parent);
93
94        tmp = gparent->rb_right;
95        if (parent != tmp) { /* parent == gparent->rb_left */
96            if (tmp && rb_is_red(tmp)) {
97                /*
98                 * Case 1 - color flips
99                 *
100                 * G g
101                 * / \ / \
102                 * p u --> P U
103                 * / /
104                 * n N
105                 *
106                 * However, since g's parent might be red, and
107                 * 4) does not allow this, we need to recurse
108                 * at g.
109                 */
110                rb_set_parent_color(tmp, gparent, RB_BLACK);
111                rb_set_parent_color(parent, gparent, RB_BLACK);
112                node = gparent;
113                parent = rb_parent(node);
114                rb_set_parent_color(node, parent, RB_RED);
115                continue;
116            }
117
118            tmp = parent->rb_right;
119            if (node == tmp) {
120                /*
121                 * Case 2 - left rotate at parent
122                 *
123                 * G G
124                 * / \ / \
125                 * p U --> n U
126                 * \ /
127                 * n p
128                 *
129                 * This still leaves us in violation of 4), the
130                 * continuation into Case 3 will fix that.
131                 */
132                parent->rb_right = tmp = node->rb_left;
133                node->rb_left = parent;
134                if (tmp)
135                    rb_set_parent_color(tmp, parent,
136                                RB_BLACK);
137                rb_set_parent_color(parent, node, RB_RED);
138                augment_rotate(parent, node);
139                parent = node;
140                tmp = node->rb_right;
141            }
142
143            /*
144             * Case 3 - right rotate at gparent
145             *
146             * G P
147             * / \ / \
148             * p U --> n g
149             * / \
150             * n U
151             */
152            gparent->rb_left = tmp; /* == parent->rb_right */
153            parent->rb_right = gparent;
154            if (tmp)
155                rb_set_parent_color(tmp, gparent, RB_BLACK);
156            __rb_rotate_set_parents(gparent, parent, root, RB_RED);
157            augment_rotate(gparent, parent);
158            break;
159        } else {
160            tmp = gparent->rb_left;
161            if (tmp && rb_is_red(tmp)) {
162                /* Case 1 - color flips */
163                rb_set_parent_color(tmp, gparent, RB_BLACK);
164                rb_set_parent_color(parent, gparent, RB_BLACK);
165                node = gparent;
166                parent = rb_parent(node);
167                rb_set_parent_color(node, parent, RB_RED);
168                continue;
169            }
170
171            tmp = parent->rb_left;
172            if (node == tmp) {
173                /* Case 2 - right rotate at parent */
174                parent->rb_left = tmp = node->rb_right;
175                node->rb_right = parent;
176                if (tmp)
177                    rb_set_parent_color(tmp, parent,
178                                RB_BLACK);
179                rb_set_parent_color(parent, node, RB_RED);
180                augment_rotate(parent, node);
181                parent = node;
182                tmp = node->rb_left;
183            }
184
185            /* Case 3 - left rotate at gparent */
186            gparent->rb_right = tmp; /* == parent->rb_left */
187            parent->rb_left = gparent;
188            if (tmp)
189                rb_set_parent_color(tmp, gparent, RB_BLACK);
190            __rb_rotate_set_parents(gparent, parent, root, RB_RED);
191            augment_rotate(gparent, parent);
192            break;
193        }
194    }
195}
196
197/*
198 * Inline version for rb_erase() use - we want to be able to inline
199 * and eliminate the dummy_rotate callback there
200 */
201static __always_inline void
202____rb_erase_color(struct rb_node *parent, struct rb_root *root,
203    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
204{
205    struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
206
207    while (true) {
208        /*
209         * Loop invariants:
210         * - node is black (or NULL on first iteration)
211         * - node is not the root (parent is not NULL)
212         * - All leaf paths going through parent and node have a
213         * black node count that is 1 lower than other leaf paths.
214         */
215        sibling = parent->rb_right;
216        if (node != sibling) { /* node == parent->rb_left */
217            if (rb_is_red(sibling)) {
218                /*
219                 * Case 1 - left rotate at parent
220                 *
221                 * P S
222                 * / \ / \
223                 * N s --> p Sr
224                 * / \ / \
225                 * Sl Sr N Sl
226                 */
227                parent->rb_right = tmp1 = sibling->rb_left;
228                sibling->rb_left = parent;
229                rb_set_parent_color(tmp1, parent, RB_BLACK);
230                __rb_rotate_set_parents(parent, sibling, root,
231                            RB_RED);
232                augment_rotate(parent, sibling);
233                sibling = tmp1;
234            }
235            tmp1 = sibling->rb_right;
236            if (!tmp1 || rb_is_black(tmp1)) {
237                tmp2 = sibling->rb_left;
238                if (!tmp2 || rb_is_black(tmp2)) {
239                    /*
240                     * Case 2 - sibling color flip
241                     * (p could be either color here)
242                     *
243                     * (p) (p)
244                     * / \ / \
245                     * N S --> N s
246                     * / \ / \
247                     * Sl Sr Sl Sr
248                     *
249                     * This leaves us violating 5) which
250                     * can be fixed by flipping p to black
251                     * if it was red, or by recursing at p.
252                     * p is red when coming from Case 1.
253                     */
254                    rb_set_parent_color(sibling, parent,
255                                RB_RED);
256                    if (rb_is_red(parent))
257                        rb_set_black(parent);
258                    else {
259                        node = parent;
260                        parent = rb_parent(node);
261                        if (parent)
262                            continue;
263                    }
264                    break;
265                }
266                /*
267                 * Case 3 - right rotate at sibling
268                 * (p could be either color here)
269                 *
270                 * (p) (p)
271                 * / \ / \
272                 * N S --> N Sl
273                 * / \ \
274                 * sl Sr s
275                 * \
276                 * Sr
277                 */
278                sibling->rb_left = tmp1 = tmp2->rb_right;
279                tmp2->rb_right = sibling;
280                parent->rb_right = tmp2;
281                if (tmp1)
282                    rb_set_parent_color(tmp1, sibling,
283                                RB_BLACK);
284                augment_rotate(sibling, tmp2);
285                tmp1 = sibling;
286                sibling = tmp2;
287            }
288            /*
289             * Case 4 - left rotate at parent + color flips
290             * (p and sl could be either color here.
291             * After rotation, p becomes black, s acquires
292             * p's color, and sl keeps its color)
293             *
294             * (p) (s)
295             * / \ / \
296             * N S --> P Sr
297             * / \ / \
298             * (sl) sr N (sl)
299             */
300            parent->rb_right = tmp2 = sibling->rb_left;
301            sibling->rb_left = parent;
302            rb_set_parent_color(tmp1, sibling, RB_BLACK);
303            if (tmp2)
304                rb_set_parent(tmp2, parent);
305            __rb_rotate_set_parents(parent, sibling, root,
306                        RB_BLACK);
307            augment_rotate(parent, sibling);
308            break;
309        } else {
310            sibling = parent->rb_left;
311            if (rb_is_red(sibling)) {
312                /* Case 1 - right rotate at parent */
313                parent->rb_left = tmp1 = sibling->rb_right;
314                sibling->rb_right = parent;
315                rb_set_parent_color(tmp1, parent, RB_BLACK);
316                __rb_rotate_set_parents(parent, sibling, root,
317                            RB_RED);
318                augment_rotate(parent, sibling);
319                sibling = tmp1;
320            }
321            tmp1 = sibling->rb_left;
322            if (!tmp1 || rb_is_black(tmp1)) {
323                tmp2 = sibling->rb_right;
324                if (!tmp2 || rb_is_black(tmp2)) {
325                    /* Case 2 - sibling color flip */
326                    rb_set_parent_color(sibling, parent,
327                                RB_RED);
328                    if (rb_is_red(parent))
329                        rb_set_black(parent);
330                    else {
331                        node = parent;
332                        parent = rb_parent(node);
333                        if (parent)
334                            continue;
335                    }
336                    break;
337                }
338                /* Case 3 - right rotate at sibling */
339                sibling->rb_right = tmp1 = tmp2->rb_left;
340                tmp2->rb_left = sibling;
341                parent->rb_left = tmp2;
342                if (tmp1)
343                    rb_set_parent_color(tmp1, sibling,
344                                RB_BLACK);
345                augment_rotate(sibling, tmp2);
346                tmp1 = sibling;
347                sibling = tmp2;
348            }
349            /* Case 4 - left rotate at parent + color flips */
350            parent->rb_left = tmp2 = sibling->rb_right;
351            sibling->rb_right = parent;
352            rb_set_parent_color(tmp1, sibling, RB_BLACK);
353            if (tmp2)
354                rb_set_parent(tmp2, parent);
355            __rb_rotate_set_parents(parent, sibling, root,
356                        RB_BLACK);
357            augment_rotate(parent, sibling);
358            break;
359        }
360    }
361}
362
363/* Non-inline version for rb_erase_augmented() use */
364void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
365    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
366{
367    ____rb_erase_color(parent, root, augment_rotate);
368}
369EXPORT_SYMBOL(__rb_erase_color);
370
371/*
372 * Non-augmented rbtree manipulation functions.
373 *
374 * We use dummy augmented callbacks here, and have the compiler optimize them
375 * out of the rb_insert_color() and rb_erase() function definitions.
376 */
377
378static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
379static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
380static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
381
382static const struct rb_augment_callbacks dummy_callbacks = {
383    dummy_propagate, dummy_copy, dummy_rotate
384};
385
386void rb_insert_color(struct rb_node *node, struct rb_root *root)
387{
388    __rb_insert(node, root, dummy_rotate);
389}
390EXPORT_SYMBOL(rb_insert_color);
391
392void rb_erase(struct rb_node *node, struct rb_root *root)
393{
394    struct rb_node *rebalance;
395    rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
396    if (rebalance)
397        ____rb_erase_color(rebalance, root, dummy_rotate);
398}
399EXPORT_SYMBOL(rb_erase);
400
401/*
402 * Augmented rbtree manipulation functions.
403 *
404 * This instantiates the same __always_inline functions as in the non-augmented
405 * case, but this time with user-defined callbacks.
406 */
407
408void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
409    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
410{
411    __rb_insert(node, root, augment_rotate);
412}
413EXPORT_SYMBOL(__rb_insert_augmented);
414
415/*
416 * This function returns the first node (in sort order) of the tree.
417 */
418struct rb_node *rb_first(const struct rb_root *root)
419{
420    struct rb_node *n;
421
422    n = root->rb_node;
423    if (!n)
424        return NULL;
425    while (n->rb_left)
426        n = n->rb_left;
427    return n;
428}
429EXPORT_SYMBOL(rb_first);
430
431struct rb_node *rb_last(const struct rb_root *root)
432{
433    struct rb_node *n;
434
435    n = root->rb_node;
436    if (!n)
437        return NULL;
438    while (n->rb_right)
439        n = n->rb_right;
440    return n;
441}
442EXPORT_SYMBOL(rb_last);
443
444struct rb_node *rb_next(const struct rb_node *node)
445{
446    struct rb_node *parent;
447
448    if (RB_EMPTY_NODE(node))
449        return NULL;
450
451    /*
452     * If we have a right-hand child, go down and then left as far
453     * as we can.
454     */
455    if (node->rb_right) {
456        node = node->rb_right;
457        while (node->rb_left)
458            node=node->rb_left;
459        return (struct rb_node *)node;
460    }
461
462    /*
463     * No right-hand children. Everything down and left is smaller than us,
464     * so any 'next' node must be in the general direction of our parent.
465     * Go up the tree; any time the ancestor is a right-hand child of its
466     * parent, keep going up. First time it's a left-hand child of its
467     * parent, said parent is our 'next' node.
468     */
469    while ((parent = rb_parent(node)) && node == parent->rb_right)
470        node = parent;
471
472    return parent;
473}
474EXPORT_SYMBOL(rb_next);
475
476struct rb_node *rb_prev(const struct rb_node *node)
477{
478    struct rb_node *parent;
479
480    if (RB_EMPTY_NODE(node))
481        return NULL;
482
483    /*
484     * If we have a left-hand child, go down and then right as far
485     * as we can.
486     */
487    if (node->rb_left) {
488        node = node->rb_left;
489        while (node->rb_right)
490            node=node->rb_right;
491        return (struct rb_node *)node;
492    }
493
494    /*
495     * No left-hand children. Go up till we find an ancestor which
496     * is a right-hand child of its parent.
497     */
498    while ((parent = rb_parent(node)) && node == parent->rb_left)
499        node = parent;
500
501    return parent;
502}
503EXPORT_SYMBOL(rb_prev);
504
505void rb_replace_node(struct rb_node *victim, struct rb_node *new,
506             struct rb_root *root)
507{
508    struct rb_node *parent = rb_parent(victim);
509
510    /* Set the surrounding nodes to point to the replacement */
511    __rb_change_child(victim, new, parent, root);
512    if (victim->rb_left)
513        rb_set_parent(victim->rb_left, new);
514    if (victim->rb_right)
515        rb_set_parent(victim->rb_right, new);
516
517    /* Copy the pointers/colour from the victim to the replacement */
518    *new = *victim;
519}
520EXPORT_SYMBOL(rb_replace_node);
521
522static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
523{
524    for (;;) {
525        if (node->rb_left)
526            node = node->rb_left;
527        else if (node->rb_right)
528            node = node->rb_right;
529        else
530            return (struct rb_node *)node;
531    }
532}
533
534struct rb_node *rb_next_postorder(const struct rb_node *node)
535{
536    const struct rb_node *parent;
537    if (!node)
538        return NULL;
539    parent = rb_parent(node);
540
541    /* If we're sitting on node, we've already seen our children */
542    if (parent && node == parent->rb_left && parent->rb_right) {
543        /* If we are the parent's left node, go to the parent's right
544         * node then all the way down to the left */
545        return rb_left_deepest_node(parent->rb_right);
546    } else
547        /* Otherwise we are the parent's right node, and the parent
548         * should be next */
549        return (struct rb_node *)parent;
550}
551EXPORT_SYMBOL(rb_next_postorder);
552
553struct rb_node *rb_first_postorder(const struct rb_root *root)
554{
555    if (!root->rb_node)
556        return NULL;
557
558    return rb_left_deepest_node(root->rb_node);
559}
560EXPORT_SYMBOL(rb_first_postorder);
561

Archive Download this file



interactive