?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
rbtree: place easiest case first in rb_erase()
[linux-2.6/libata-dev.git] / lib / rbtree.c
blobbde1b5c5fb33d4f45266c844b62412a90d86baec
1 /*
2 Red Black Trees
3 (C) 1999 Andrea Arcangeli <andrea@suse.de>
4 (C) 2002 David Woodhouse <dwmw2@infradead.org>
5
6 This program is free software; you can redistribute it and/or modify
7 it under the terms of the GNU General Public License as published by
8 the Free Software Foundation; either version 2 of the License, or
9 (at your option) any later version.
10
11 This program is distributed in the hope that it will be useful,
12 but WITHOUT ANY WARRANTY; without even the implied warranty of
13 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 GNU General Public License for more details.
15
16 You should have received a copy of the GNU General Public License
17 along with this program; if not, write to the Free Software
18 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
19
20 linux/lib/rbtree.c
21 */
22
23 #include <linux/rbtree.h>
24 #include <linux/export.h>
25
26 /*
27 * red-black trees properties: http://en.wikipedia.org/wiki/Rbtree
28 *
29 * 1) A node is either red or black
30 * 2) The root is black
31 * 3) All leaves (NULL) are black
32 * 4) Both children of every red node are black
33 * 5) Every simple path from root to leaves contains the same number
34 * of black nodes.
35 *
36 * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
37 * consecutive red nodes in a path and every red node is therefore followed by
38 * a black. So if B is the number of black nodes on every simple path (as per
39 * 5), then the longest possible path due to 4 is 2B.
40 *
41 * We shall indicate color with case, where black nodes are uppercase and red
42 * nodes will be lowercase. Unknown color nodes shall be drawn as red within
43 * parentheses and have some accompanying text comment.
44 */
45
46 #define RB_RED 0
47 #define RB_BLACK 1
48
49 #define rb_color(r) ((r)->__rb_parent_color & 1)
50 #define rb_is_red(r) (!rb_color(r))
51 #define rb_is_black(r) rb_color(r)
52
53 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
54 {
55 rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
56 }
57
58 static inline void rb_set_parent_color(struct rb_node *rb,
59 struct rb_node *p, int color)
60 {
61 rb->__rb_parent_color = (unsigned long)p | color;
62 }
63
64 static inline struct rb_node *rb_red_parent(struct rb_node *red)
65 {
66 return (struct rb_node *)red->__rb_parent_color;
67 }
68
69 static inline void
70 __rb_change_child(struct rb_node *old, struct rb_node *new,
71 struct rb_node *parent, struct rb_root *root)
72 {
73 if (parent) {
74 if (parent->rb_left == old)
75 parent->rb_left = new;
76 else
77 parent->rb_right = new;
78 } else
79 root->rb_node = new;
80 }
81
82 /*
83 * Helper function for rotations:
84 * - old's parent and color get assigned to new
85 * - old gets assigned new as a parent and 'color' as a color.
86 */
87 static inline void
88 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
89 struct rb_root *root, int color)
90 {
91 struct rb_node *parent = rb_parent(old);
92 new->__rb_parent_color = old->__rb_parent_color;
93 rb_set_parent_color(old, new, color);
94 __rb_change_child(old, new, parent, root);
95 }
96
97 void rb_insert_color(struct rb_node *node, struct rb_root *root)
98 {
99 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
100
101 while (true) {
102 /*
103 * Loop invariant: node is red
104 *
105 * If there is a black parent, we are done.
106 * Otherwise, take some corrective action as we don't
107 * want a red root or two consecutive red nodes.
108 */
109 if (!parent) {
110 rb_set_parent_color(node, NULL, RB_BLACK);
111 break;
112 } else if (rb_is_black(parent))
113 break;
114
115 gparent = rb_red_parent(parent);
116
117 tmp = gparent->rb_right;
118 if (parent != tmp) { /* parent == gparent->rb_left */
119 if (tmp && rb_is_red(tmp)) {
120 /*
121 * Case 1 - color flips
122 *
123 * G g
124 * / \ / \
125 * p u --> P U
126 * / /
127 * n N
128 *
129 * However, since g's parent might be red, and
130 * 4) does not allow this, we need to recurse
131 * at g.
132 */
133 rb_set_parent_color(tmp, gparent, RB_BLACK);
134 rb_set_parent_color(parent, gparent, RB_BLACK);
135 node = gparent;
136 parent = rb_parent(node);
137 rb_set_parent_color(node, parent, RB_RED);
138 continue;
139 }
140
141 tmp = parent->rb_right;
142 if (node == tmp) {
143 /*
144 * Case 2 - left rotate at parent
145 *
146 * G G
147 * / \ / \
148 * p U --> n U
149 * \ /
150 * n p
151 *
152 * This still leaves us in violation of 4), the
153 * continuation into Case 3 will fix that.
154 */
155 parent->rb_right = tmp = node->rb_left;
156 node->rb_left = parent;
157 if (tmp)
158 rb_set_parent_color(tmp, parent,
159 RB_BLACK);
160 rb_set_parent_color(parent, node, RB_RED);
161 parent = node;
162 tmp = node->rb_right;
163 }
164
165 /*
166 * Case 3 - right rotate at gparent
167 *
168 * G P
169 * / \ / \
170 * p U --> n g
171 * / \
172 * n U
173 */
174 gparent->rb_left = tmp; /* == parent->rb_right */
175 parent->rb_right = gparent;
176 if (tmp)
177 rb_set_parent_color(tmp, gparent, RB_BLACK);
178 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
179 break;
180 } else {
181 tmp = gparent->rb_left;
182 if (tmp && rb_is_red(tmp)) {
183 /* Case 1 - color flips */
184 rb_set_parent_color(tmp, gparent, RB_BLACK);
185 rb_set_parent_color(parent, gparent, RB_BLACK);
186 node = gparent;
187 parent = rb_parent(node);
188 rb_set_parent_color(node, parent, RB_RED);
189 continue;
190 }
191
192 tmp = parent->rb_left;
193 if (node == tmp) {
194 /* Case 2 - right rotate at parent */
195 parent->rb_left = tmp = node->rb_right;
196 node->rb_right = parent;
197 if (tmp)
198 rb_set_parent_color(tmp, parent,
199 RB_BLACK);
200 rb_set_parent_color(parent, node, RB_RED);
201 parent = node;
202 tmp = node->rb_left;
203 }
204
205 /* Case 3 - left rotate at gparent */
206 gparent->rb_right = tmp; /* == parent->rb_left */
207 parent->rb_left = gparent;
208 if (tmp)
209 rb_set_parent_color(tmp, gparent, RB_BLACK);
210 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
211 break;
212 }
213 }
214 }
215 EXPORT_SYMBOL(rb_insert_color);
216
217 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
218 struct rb_root *root)
219 {
220 struct rb_node *sibling, *tmp1, *tmp2;
221
222 while (true) {
223 /*
224 * Loop invariant: all leaf paths going through node have a
225 * black node count that is 1 lower than other leaf paths.
226 *
227 * If node is red, we can flip it to black to adjust.
228 * If node is the root, all leaf paths go through it.
229 * Otherwise, we need to adjust the tree through color flips
230 * and tree rotations as per one of the 4 cases below.
231 */
232 if (node && rb_is_red(node)) {
233 rb_set_parent_color(node, parent, RB_BLACK);
234 break;
235 } else if (!parent) {
236 break;
237 }
238 sibling = parent->rb_right;
239 if (node != sibling) { /* node == parent->rb_left */
240 if (rb_is_red(sibling)) {
241 /*
242 * Case 1 - left rotate at parent
243 *
244 * P S
245 * / \ / \
246 * N s --> p Sr
247 * / \ / \
248 * Sl Sr N Sl
249 */
250 parent->rb_right = tmp1 = sibling->rb_left;
251 sibling->rb_left = parent;
252 rb_set_parent_color(tmp1, parent, RB_BLACK);
253 __rb_rotate_set_parents(parent, sibling, root,
254 RB_RED);
255 sibling = tmp1;
256 }
257 tmp1 = sibling->rb_right;
258 if (!tmp1 || rb_is_black(tmp1)) {
259 tmp2 = sibling->rb_left;
260 if (!tmp2 || rb_is_black(tmp2)) {
261 /*
262 * Case 2 - sibling color flip
263 * (p could be either color here)
264 *
265 * (p) (p)
266 * / \ / \
267 * N S --> N s
268 * / \ / \
269 * Sl Sr Sl Sr
270 *
271 * This leaves us violating 5), so
272 * recurse at p. If p is red, the
273 * recursion will just flip it to black
274 * and exit. If coming from Case 1,
275 * p is known to be red.
276 */
277 rb_set_parent_color(sibling, parent,
278 RB_RED);
279 node = parent;
280 parent = rb_parent(node);
281 continue;
282 }
283 /*
284 * Case 3 - right rotate at sibling
285 * (p could be either color here)
286 *
287 * (p) (p)
288 * / \ / \
289 * N S --> N Sl
290 * / \ \
291 * sl Sr s
292 * \
293 * Sr
294 */
295 sibling->rb_left = tmp1 = tmp2->rb_right;
296 tmp2->rb_right = sibling;
297 parent->rb_right = tmp2;
298 if (tmp1)
299 rb_set_parent_color(tmp1, sibling,
300 RB_BLACK);
301 tmp1 = sibling;
302 sibling = tmp2;
303 }
304 /*
305 * Case 4 - left rotate at parent + color flips
306 * (p and sl could be either color here.
307 * After rotation, p becomes black, s acquires
308 * p's color, and sl keeps its color)
309 *
310 * (p) (s)
311 * / \ / \
312 * N S --> P Sr
313 * / \ / \
314 * (sl) sr N (sl)
315 */
316 parent->rb_right = tmp2 = sibling->rb_left;
317 sibling->rb_left = parent;
318 rb_set_parent_color(tmp1, sibling, RB_BLACK);
319 if (tmp2)
320 rb_set_parent(tmp2, parent);
321 __rb_rotate_set_parents(parent, sibling, root,
322 RB_BLACK);
323 break;
324 } else {
325 sibling = parent->rb_left;
326 if (rb_is_red(sibling)) {
327 /* Case 1 - right rotate at parent */
328 parent->rb_left = tmp1 = sibling->rb_right;
329 sibling->rb_right = parent;
330 rb_set_parent_color(tmp1, parent, RB_BLACK);
331 __rb_rotate_set_parents(parent, sibling, root,
332 RB_RED);
333 sibling = tmp1;
334 }
335 tmp1 = sibling->rb_left;
336 if (!tmp1 || rb_is_black(tmp1)) {
337 tmp2 = sibling->rb_right;
338 if (!tmp2 || rb_is_black(tmp2)) {
339 /* Case 2 - sibling color flip */
340 rb_set_parent_color(sibling, parent,
341 RB_RED);
342 node = parent;
343 parent = rb_parent(node);
344 continue;
345 }
346 /* Case 3 - right rotate at sibling */
347 sibling->rb_right = tmp1 = tmp2->rb_left;
348 tmp2->rb_left = sibling;
349 parent->rb_left = tmp2;
350 if (tmp1)
351 rb_set_parent_color(tmp1, sibling,
352 RB_BLACK);
353 tmp1 = sibling;
354 sibling = tmp2;
355 }
356 /* Case 4 - left rotate at parent + color flips */
357 parent->rb_left = tmp2 = sibling->rb_right;
358 sibling->rb_right = parent;
359 rb_set_parent_color(tmp1, sibling, RB_BLACK);
360 if (tmp2)
361 rb_set_parent(tmp2, parent);
362 __rb_rotate_set_parents(parent, sibling, root,
363 RB_BLACK);
364 break;
365 }
366 }
367 }
368
369 void rb_erase(struct rb_node *node, struct rb_root *root)
370 {
371 struct rb_node *child = node->rb_right, *tmp = node->rb_left;
372 struct rb_node *parent;
373 int color;
374
375 if (!tmp) {
376 case1:
377 /* Case 1: node to erase has no more than 1 child (easy!) */
378
379 parent = rb_parent(node);
380 color = rb_color(node);
381
382 if (child)
383 rb_set_parent(child, parent);
384 __rb_change_child(node, child, parent, root);
385 } else if (!child) {
386 /* Still case 1, but this time the child is node->rb_left */
387 child = tmp;
388 goto case1;
389 } else {
390 struct rb_node *old = node, *left;
391
392 node = child;
393 while ((left = node->rb_left) != NULL)
394 node = left;
395
396 __rb_change_child(old, node, rb_parent(old), root);
397
398 child = node->rb_right;
399 parent = rb_parent(node);
400 color = rb_color(node);
401
402 if (parent == old) {
403 parent = node;
404 } else {
405 if (child)
406 rb_set_parent(child, parent);
407 parent->rb_left = child;
408
409 node->rb_right = old->rb_right;
410 rb_set_parent(old->rb_right, node);
411 }
412
413 node->__rb_parent_color = old->__rb_parent_color;
414 node->rb_left = old->rb_left;
415 rb_set_parent(old->rb_left, node);
416 }
417
418 if (color == RB_BLACK)
419 __rb_erase_color(child, parent, root);
420 }
421 EXPORT_SYMBOL(rb_erase);
422
423 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
424 {
425 struct rb_node *parent;
426
427 up:
428 func(node, data);
429 parent = rb_parent(node);
430 if (!parent)
431 return;
432
433 if (node == parent->rb_left && parent->rb_right)
434 func(parent->rb_right, data);
435 else if (parent->rb_left)
436 func(parent->rb_left, data);
437
438 node = parent;
439 goto up;
440 }
441
442 /*
443 * after inserting @node into the tree, update the tree to account for
444 * both the new entry and any damage done by rebalance
445 */
446 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
447 {
448 if (node->rb_left)
449 node = node->rb_left;
450 else if (node->rb_right)
451 node = node->rb_right;
452
453 rb_augment_path(node, func, data);
454 }
455 EXPORT_SYMBOL(rb_augment_insert);
456
457 /*
458 * before removing the node, find the deepest node on the rebalance path
459 * that will still be there after @node gets removed
460 */
461 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
462 {
463 struct rb_node *deepest;
464
465 if (!node->rb_right && !node->rb_left)
466 deepest = rb_parent(node);
467 else if (!node->rb_right)
468 deepest = node->rb_left;
469 else if (!node->rb_left)
470 deepest = node->rb_right;
471 else {
472 deepest = rb_next(node);
473 if (deepest->rb_right)
474 deepest = deepest->rb_right;
475 else if (rb_parent(deepest) != node)
476 deepest = rb_parent(deepest);
477 }
478
479 return deepest;
480 }
481 EXPORT_SYMBOL(rb_augment_erase_begin);
482
483 /*
484 * after removal, update the tree to account for the removed entry
485 * and any rebalance damage.
486 */
487 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
488 {
489 if (node)
490 rb_augment_path(node, func, data);
491 }
492 EXPORT_SYMBOL(rb_augment_erase_end);
493
494 /*
495 * This function returns the first node (in sort order) of the tree.
496 */
497 struct rb_node *rb_first(const struct rb_root *root)
498 {
499 struct rb_node *n;
500
501 n = root->rb_node;
502 if (!n)
503 return NULL;
504 while (n->rb_left)
505 n = n->rb_left;
506 return n;
507 }
508 EXPORT_SYMBOL(rb_first);
509
510 struct rb_node *rb_last(const struct rb_root *root)
511 {
512 struct rb_node *n;
513
514 n = root->rb_node;
515 if (!n)
516 return NULL;
517 while (n->rb_right)
518 n = n->rb_right;
519 return n;
520 }
521 EXPORT_SYMBOL(rb_last);
522
523 struct rb_node *rb_next(const struct rb_node *node)
524 {
525 struct rb_node *parent;
526
527 if (RB_EMPTY_NODE(node))
528 return NULL;
529
530 /*
531 * If we have a right-hand child, go down and then left as far
532 * as we can.
533 */
534 if (node->rb_right) {
535 node = node->rb_right;
536 while (node->rb_left)
537 node=node->rb_left;
538 return (struct rb_node *)node;
539 }
540
541 /*
542 * No right-hand children. Everything down and left is smaller than us,
543 * so any 'next' node must be in the general direction of our parent.
544 * Go up the tree; any time the ancestor is a right-hand child of its
545 * parent, keep going up. First time it's a left-hand child of its
546 * parent, said parent is our 'next' node.
547 */
548 while ((parent = rb_parent(node)) && node == parent->rb_right)
549 node = parent;
550
551 return parent;
552 }
553 EXPORT_SYMBOL(rb_next);
554
555 struct rb_node *rb_prev(const struct rb_node *node)
556 {
557 struct rb_node *parent;
558
559 if (RB_EMPTY_NODE(node))
560 return NULL;
561
562 /*
563 * If we have a left-hand child, go down and then right as far
564 * as we can.
565 */
566 if (node->rb_left) {
567 node = node->rb_left;
568 while (node->rb_right)
569 node=node->rb_right;
570 return (struct rb_node *)node;
571 }
572
573 /*
574 * No left-hand children. Go up till we find an ancestor which
575 * is a right-hand child of its parent.
576 */
577 while ((parent = rb_parent(node)) && node == parent->rb_left)
578 node = parent;
579
580 return parent;
581 }
582 EXPORT_SYMBOL(rb_prev);
583
584 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
585 struct rb_root *root)
586 {
587 struct rb_node *parent = rb_parent(victim);
588
589 /* Set the surrounding nodes to point to the replacement */
590 __rb_change_child(victim, new, parent, root);
591 if (victim->rb_left)
592 rb_set_parent(victim->rb_left, new);
593 if (victim->rb_right)
594 rb_set_parent(victim->rb_right, new);
595
596 /* Copy the pointers/colour from the victim to the replacement */
597 *new = *victim;
598 }
599 EXPORT_SYMBOL(rb_replace_node);
Linux 2.6 libata-dev (mirror)