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.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 #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 59 static inline void rb_set_black(struct rb_node *rb) 60 { 61 rb->__rb_parent_color |= RB_BLACK; 62 } 63 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 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 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 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 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 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 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 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 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data) 473 { 474 struct rb_node *parent; 475 476 up: 477 func(node, data); 478 parent = rb_parent(node); 479 if (!parent) 480 return; 481 482 if (node == parent->rb_left && parent->rb_right) 483 func(parent->rb_right, data); 484 else if (parent->rb_left) 485 func(parent->rb_left, data); 486 487 node = parent; 488 goto up; 489 } 490 491 /* 492 * after inserting @node into the tree, update the tree to account for 493 * both the new entry and any damage done by rebalance 494 */ 495 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data) 496 { 497 if (node->rb_left) 498 node = node->rb_left; 499 else if (node->rb_right) 500 node = node->rb_right; 501 502 rb_augment_path(node, func, data); 503 } 504 EXPORT_SYMBOL(rb_augment_insert); 505 506 /* 507 * before removing the node, find the deepest node on the rebalance path 508 * that will still be there after @node gets removed 509 */ 510 struct rb_node *rb_augment_erase_begin(struct rb_node *node) 511 { 512 struct rb_node *deepest; 513 514 if (!node->rb_right && !node->rb_left) 515 deepest = rb_parent(node); 516 else if (!node->rb_right) 517 deepest = node->rb_left; 518 else if (!node->rb_left) 519 deepest = node->rb_right; 520 else { 521 deepest = rb_next(node); 522 if (deepest->rb_right) 523 deepest = deepest->rb_right; 524 else if (rb_parent(deepest) != node) 525 deepest = rb_parent(deepest); 526 } 527 528 return deepest; 529 } 530 EXPORT_SYMBOL(rb_augment_erase_begin); 531 532 /* 533 * after removal, update the tree to account for the removed entry 534 * and any rebalance damage. 535 */ 536 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data) 537 { 538 if (node) 539 rb_augment_path(node, func, data); 540 } 541 EXPORT_SYMBOL(rb_augment_erase_end); 542 543 /* 544 * This function returns the first node (in sort order) of the tree. 545 */ 546 struct rb_node *rb_first(const struct rb_root *root) 547 { 548 struct rb_node *n; 549 550 n = root->rb_node; 551 if (!n) 552 return NULL; 553 while (n->rb_left) 554 n = n->rb_left; 555 return n; 556 } 557 EXPORT_SYMBOL(rb_first); 558 559 struct rb_node *rb_last(const struct rb_root *root) 560 { 561 struct rb_node *n; 562 563 n = root->rb_node; 564 if (!n) 565 return NULL; 566 while (n->rb_right) 567 n = n->rb_right; 568 return n; 569 } 570 EXPORT_SYMBOL(rb_last); 571 572 struct rb_node *rb_next(const struct rb_node *node) 573 { 574 struct rb_node *parent; 575 576 if (RB_EMPTY_NODE(node)) 577 return NULL; 578 579 /* 580 * If we have a right-hand child, go down and then left as far 581 * as we can. 582 */ 583 if (node->rb_right) { 584 node = node->rb_right; 585 while (node->rb_left) 586 node=node->rb_left; 587 return (struct rb_node *)node; 588 } 589 590 /* 591 * No right-hand children. Everything down and left is smaller than us, 592 * so any 'next' node must be in the general direction of our parent. 593 * Go up the tree; any time the ancestor is a right-hand child of its 594 * parent, keep going up. First time it's a left-hand child of its 595 * parent, said parent is our 'next' node. 596 */ 597 while ((parent = rb_parent(node)) && node == parent->rb_right) 598 node = parent; 599 600 return parent; 601 } 602 EXPORT_SYMBOL(rb_next); 603 604 struct rb_node *rb_prev(const struct rb_node *node) 605 { 606 struct rb_node *parent; 607 608 if (RB_EMPTY_NODE(node)) 609 return NULL; 610 611 /* 612 * If we have a left-hand child, go down and then right as far 613 * as we can. 614 */ 615 if (node->rb_left) { 616 node = node->rb_left; 617 while (node->rb_right) 618 node=node->rb_right; 619 return (struct rb_node *)node; 620 } 621 622 /* 623 * No left-hand children. Go up till we find an ancestor which 624 * is a right-hand child of its parent. 625 */ 626 while ((parent = rb_parent(node)) && node == parent->rb_left) 627 node = parent; 628 629 return parent; 630 } 631 EXPORT_SYMBOL(rb_prev); 632 633 void rb_replace_node(struct rb_node *victim, struct rb_node *new, 634 struct rb_root *root) 635 { 636 struct rb_node *parent = rb_parent(victim); 637 638 /* Set the surrounding nodes to point to the replacement */ 639 __rb_change_child(victim, new, parent, root); 640 if (victim->rb_left) 641 rb_set_parent(victim->rb_left, new); 642 if (victim->rb_right) 643 rb_set_parent(victim->rb_right, new); 644 645 /* Copy the pointers/colour from the victim to the replacement */ 646 *new = *victim; 647 } 648 EXPORT_SYMBOL(rb_replace_node); 649