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