Line data Source code
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 :
47 : static inline void rb_set_black(struct rb_node *rb)
48 : {
49 85992 : rb->__rb_parent_color |= RB_BLACK;
50 : }
51 :
52 : static inline struct rb_node *rb_red_parent(struct rb_node *red)
53 : {
54 4111592 : 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 : */
62 : static inline void
63 : __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
64 : struct rb_root *root, int color)
65 : {
66 655951 : struct rb_node *parent = rb_parent(old);
67 655951 : 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 :
72 : static __always_inline void
73 3257836 : __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 3560308 : if (!parent) {
87 : rb_set_parent_color(node, NULL, RB_BLACK);
88 : break;
89 2702964 : } else if (rb_is_black(parent))
90 : break;
91 :
92 : gparent = rb_red_parent(parent);
93 :
94 853756 : tmp = gparent->rb_right;
95 853756 : if (parent != tmp) { /* parent == gparent->rb_left */
96 342170 : 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 37202 : parent = rb_parent(node);
114 : rb_set_parent_color(node, parent, RB_RED);
115 37202 : continue;
116 : }
117 :
118 304968 : tmp = parent->rb_right;
119 304968 : 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 281385 : parent->rb_right = tmp = node->rb_left;
133 281385 : node->rb_left = parent;
134 281385 : if (tmp)
135 : rb_set_parent_color(tmp, parent,
136 : RB_BLACK);
137 : rb_set_parent_color(parent, node, RB_RED);
138 31845 : augment_rotate(parent, node);
139 : parent = node;
140 281385 : 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 304968 : gparent->rb_left = tmp; /* == parent->rb_right */
153 304968 : parent->rb_right = gparent;
154 304968 : if (tmp)
155 : rb_set_parent_color(tmp, gparent, RB_BLACK);
156 : __rb_rotate_set_parents(gparent, parent, root, RB_RED);
157 46179 : augment_rotate(gparent, parent);
158 : break;
159 : } else {
160 511586 : tmp = gparent->rb_left;
161 511586 : 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 265270 : parent = rb_parent(node);
167 : rb_set_parent_color(node, parent, RB_RED);
168 265270 : continue;
169 : }
170 :
171 246316 : tmp = parent->rb_left;
172 246316 : if (node == tmp) {
173 : /* Case 2 - right rotate at parent */
174 29440 : parent->rb_left = tmp = node->rb_right;
175 29440 : node->rb_right = parent;
176 29440 : if (tmp)
177 : rb_set_parent_color(tmp, parent,
178 : RB_BLACK);
179 : rb_set_parent_color(parent, node, RB_RED);
180 13367 : augment_rotate(parent, node);
181 : parent = node;
182 29440 : tmp = node->rb_left;
183 : }
184 :
185 : /* Case 3 - left rotate at gparent */
186 246316 : gparent->rb_right = tmp; /* == parent->rb_left */
187 246316 : parent->rb_left = gparent;
188 246316 : if (tmp)
189 : rb_set_parent_color(tmp, gparent, RB_BLACK);
190 : __rb_rotate_set_parents(gparent, parent, root, RB_RED);
191 108221 : 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 : */
201 : static __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 250813 : sibling = parent->rb_right;
216 250813 : if (node != sibling) { /* node == parent->rb_left */
217 184632 : 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 51423 : parent->rb_right = tmp1 = sibling->rb_left;
228 51423 : 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 765 : augment_rotate(parent, sibling);
233 : sibling = tmp1;
234 : }
235 184619 : tmp1 = sibling->rb_right;
236 184619 : if (!tmp1 || rb_is_black(tmp1)) {
237 163389 : tmp2 = sibling->rb_left;
238 163389 : 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 151559 : if (rb_is_red(parent))
257 : rb_set_black(parent);
258 : else {
259 : node = parent;
260 89934 : parent = rb_parent(node);
261 89934 : if (parent)
262 64117 : 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 11830 : sibling->rb_left = tmp1 = tmp2->rb_right;
279 11830 : tmp2->rb_right = sibling;
280 11830 : parent->rb_right = tmp2;
281 11830 : if (tmp1)
282 : rb_set_parent_color(tmp1, sibling,
283 : RB_BLACK);
284 2520 : 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 33060 : parent->rb_right = tmp2 = sibling->rb_left;
301 33060 : sibling->rb_left = parent;
302 : rb_set_parent_color(tmp1, sibling, RB_BLACK);
303 33060 : if (tmp2)
304 : rb_set_parent(tmp2, parent);
305 : __rb_rotate_set_parents(parent, sibling, root,
306 : RB_BLACK);
307 7239 : augment_rotate(parent, sibling);
308 : break;
309 : } else {
310 66181 : sibling = parent->rb_left;
311 66181 : if (rb_is_red(sibling)) {
312 : /* Case 1 - right rotate at parent */
313 7000 : parent->rb_left = tmp1 = sibling->rb_right;
314 7000 : 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 1764 : augment_rotate(parent, sibling);
319 : sibling = tmp1;
320 : }
321 66181 : tmp1 = sibling->rb_left;
322 66181 : if (!tmp1 || rb_is_black(tmp1)) {
323 54893 : tmp2 = sibling->rb_right;
324 54893 : if (!tmp2 || rb_is_black(tmp2)) {
325 : /* Case 2 - sibling color flip */
326 : rb_set_parent_color(sibling, parent,
327 : RB_RED);
328 52997 : if (rb_is_red(parent))
329 : rb_set_black(parent);
330 : else {
331 : node = parent;
332 28630 : parent = rb_parent(node);
333 28630 : if (parent)
334 15226 : continue;
335 : }
336 : break;
337 : }
338 : /* Case 3 - right rotate at sibling */
339 1896 : sibling->rb_right = tmp1 = tmp2->rb_left;
340 1896 : tmp2->rb_left = sibling;
341 1896 : parent->rb_left = tmp2;
342 1896 : if (tmp1)
343 : rb_set_parent_color(tmp1, sibling,
344 : RB_BLACK);
345 1523 : augment_rotate(sibling, tmp2);
346 : tmp1 = sibling;
347 : sibling = tmp2;
348 : }
349 : /* Case 4 - left rotate at parent + color flips */
350 13184 : parent->rb_left = tmp2 = sibling->rb_right;
351 13184 : sibling->rb_right = parent;
352 : rb_set_parent_color(tmp1, sibling, RB_BLACK);
353 13184 : if (tmp2)
354 : rb_set_parent(tmp2, parent);
355 : __rb_rotate_set_parents(parent, sibling, root,
356 : RB_BLACK);
357 2538 : augment_rotate(parent, sibling);
358 : break;
359 : }
360 : }
361 : }
362 :
363 : /* Non-inline version for rb_erase_augmented() use */
364 45053 : void __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 45053 : }
369 : EXPORT_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 :
378 : static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
379 : static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
380 : static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
381 :
382 : static const struct rb_augment_callbacks dummy_callbacks = {
383 : dummy_propagate, dummy_copy, dummy_rotate
384 : };
385 :
386 2747988 : void rb_insert_color(struct rb_node *node, struct rb_root *root)
387 : {
388 : __rb_insert(node, root, dummy_rotate);
389 2747988 : }
390 : EXPORT_SYMBOL(rb_insert_color);
391 :
392 2711219 : void 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 2711219 : if (rebalance)
397 : ____rb_erase_color(rebalance, root, dummy_rotate);
398 2711219 : }
399 : EXPORT_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 :
408 509848 : void __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 509827 : }
413 : EXPORT_SYMBOL(__rb_insert_augmented);
414 :
415 : /*
416 : * This function returns the first node (in sort order) of the tree.
417 : */
418 5063 : struct rb_node *rb_first(const struct rb_root *root)
419 : {
420 : struct rb_node *n;
421 :
422 5063 : n = root->rb_node;
423 5063 : if (!n)
424 : return NULL;
425 6432 : while (n->rb_left)
426 : n = n->rb_left;
427 : return n;
428 : }
429 : EXPORT_SYMBOL(rb_first);
430 :
431 0 : struct rb_node *rb_last(const struct rb_root *root)
432 : {
433 : struct rb_node *n;
434 :
435 0 : n = root->rb_node;
436 0 : if (!n)
437 : return NULL;
438 0 : while (n->rb_right)
439 : n = n->rb_right;
440 : return n;
441 : }
442 : EXPORT_SYMBOL(rb_last);
443 :
444 1791064 : struct rb_node *rb_next(const struct rb_node *node)
445 : {
446 : struct rb_node *parent;
447 :
448 1791064 : 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 1791065 : if (node->rb_right) {
456 : node = node->rb_right;
457 535484 : 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 1281600 : while ((parent = rb_parent(node)) && node == parent->rb_right)
470 : node = parent;
471 :
472 : return parent;
473 : }
474 : EXPORT_SYMBOL(rb_next);
475 :
476 139873 : struct rb_node *rb_prev(const struct rb_node *node)
477 : {
478 : struct rb_node *parent;
479 :
480 139873 : 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 139873 : if (node->rb_left) {
488 : node = node->rb_left;
489 26647 : 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 142437 : while ((parent = rb_parent(node)) && node == parent->rb_left)
499 : node = parent;
500 :
501 : return parent;
502 : }
503 : EXPORT_SYMBOL(rb_prev);
504 :
505 0 : void rb_replace_node(struct rb_node *victim, struct rb_node *new,
506 : struct rb_root *root)
507 : {
508 0 : 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 0 : if (victim->rb_left)
513 : rb_set_parent(victim->rb_left, new);
514 0 : 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 0 : *new = *victim;
519 0 : }
520 : EXPORT_SYMBOL(rb_replace_node);
521 :
522 : static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
523 : {
524 : for (;;) {
525 24013 : if (node->rb_left)
526 : node = node->rb_left;
527 12817 : else if (node->rb_right)
528 : node = node->rb_right;
529 : else
530 : return (struct rb_node *)node;
531 : }
532 : }
533 :
534 24013 : struct rb_node *rb_next_postorder(const struct rb_node *node)
535 : {
536 : const struct rb_node *parent;
537 24013 : if (!node)
538 : return NULL;
539 24013 : parent = rb_parent(node);
540 :
541 : /* If we're sitting on node, we've already seen our children */
542 24013 : 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 : }
551 : EXPORT_SYMBOL(rb_next_postorder);
552 :
553 3482 : struct rb_node *rb_first_postorder(const struct rb_root *root)
554 : {
555 3482 : if (!root->rb_node)
556 : return NULL;
557 :
558 : return rb_left_deepest_node(root->rb_node);
559 : }
560 : EXPORT_SYMBOL(rb_first_postorder);
|
