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 static __always_inline void 109 __rb_insert(struct rb_node *node, struct rb_root *root, 110 void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 111 { 112 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp; 113 114 while (true) { 115 /* 116 * Loop invariant: node is red 117 * 118 * If there is a black parent, we are done. 119 * Otherwise, take some corrective action as we don't 120 * want a red root or two consecutive red nodes. 121 */ 122 if (!parent) { 123 rb_set_parent_color(node, NULL, RB_BLACK); 124 break; 125 } else if (rb_is_black(parent)) 126 break; 127 128 gparent = rb_red_parent(parent); 129 130 tmp = gparent->rb_right; 131 if (parent != tmp) { /* parent == gparent->rb_left */ 132 if (tmp && rb_is_red(tmp)) { 133 /* 134 * Case 1 - color flips 135 * 136 * G g 137 * / \ / \ 138 * p u --> P U 139 * / / 140 * n N 141 * 142 * However, since g's parent might be red, and 143 * 4) does not allow this, we need to recurse 144 * at g. 145 */ 146 rb_set_parent_color(tmp, gparent, RB_BLACK); 147 rb_set_parent_color(parent, gparent, RB_BLACK); 148 node = gparent; 149 parent = rb_parent(node); 150 rb_set_parent_color(node, parent, RB_RED); 151 continue; 152 } 153 154 tmp = parent->rb_right; 155 if (node == tmp) { 156 /* 157 * Case 2 - left rotate at parent 158 * 159 * G G 160 * / \ / \ 161 * p U --> n U 162 * \ / 163 * n p 164 * 165 * This still leaves us in violation of 4), the 166 * continuation into Case 3 will fix that. 167 */ 168 parent->rb_right = tmp = node->rb_left; 169 node->rb_left = parent; 170 if (tmp) 171 rb_set_parent_color(tmp, parent, 172 RB_BLACK); 173 rb_set_parent_color(parent, node, RB_RED); 174 augment_rotate(parent, node); 175 parent = node; 176 tmp = node->rb_right; 177 } 178 179 /* 180 * Case 3 - right rotate at gparent 181 * 182 * G P 183 * / \ / \ 184 * p U --> n g 185 * / \ 186 * n U 187 */ 188 gparent->rb_left = tmp; /* == parent->rb_right */ 189 parent->rb_right = gparent; 190 if (tmp) 191 rb_set_parent_color(tmp, gparent, RB_BLACK); 192 __rb_rotate_set_parents(gparent, parent, root, RB_RED); 193 augment_rotate(gparent, parent); 194 break; 195 } else { 196 tmp = gparent->rb_left; 197 if (tmp && rb_is_red(tmp)) { 198 /* Case 1 - color flips */ 199 rb_set_parent_color(tmp, gparent, RB_BLACK); 200 rb_set_parent_color(parent, gparent, RB_BLACK); 201 node = gparent; 202 parent = rb_parent(node); 203 rb_set_parent_color(node, parent, RB_RED); 204 continue; 205 } 206 207 tmp = parent->rb_left; 208 if (node == tmp) { 209 /* Case 2 - right rotate at parent */ 210 parent->rb_left = tmp = node->rb_right; 211 node->rb_right = parent; 212 if (tmp) 213 rb_set_parent_color(tmp, parent, 214 RB_BLACK); 215 rb_set_parent_color(parent, node, RB_RED); 216 augment_rotate(parent, node); 217 parent = node; 218 tmp = node->rb_left; 219 } 220 221 /* Case 3 - left rotate at gparent */ 222 gparent->rb_right = tmp; /* == parent->rb_left */ 223 parent->rb_left = gparent; 224 if (tmp) 225 rb_set_parent_color(tmp, gparent, RB_BLACK); 226 __rb_rotate_set_parents(gparent, parent, root, RB_RED); 227 augment_rotate(gparent, parent); 228 break; 229 } 230 } 231 } 232 233 static __always_inline void 234 __rb_erase_color(struct rb_node *parent, struct rb_root *root, 235 const struct rb_augment_callbacks *augment) 236 { 237 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; 238 239 while (true) { 240 /* 241 * Loop invariants: 242 * - node is black (or NULL on first iteration) 243 * - node is not the root (parent is not NULL) 244 * - All leaf paths going through parent and node have a 245 * black node count that is 1 lower than other leaf paths. 246 */ 247 sibling = parent->rb_right; 248 if (node != sibling) { /* node == parent->rb_left */ 249 if (rb_is_red(sibling)) { 250 /* 251 * Case 1 - left rotate at parent 252 * 253 * P S 254 * / \ / \ 255 * N s --> p Sr 256 * / \ / \ 257 * Sl Sr N Sl 258 */ 259 parent->rb_right = tmp1 = sibling->rb_left; 260 sibling->rb_left = parent; 261 rb_set_parent_color(tmp1, parent, RB_BLACK); 262 __rb_rotate_set_parents(parent, sibling, root, 263 RB_RED); 264 augment->rotate(parent, sibling); 265 sibling = tmp1; 266 } 267 tmp1 = sibling->rb_right; 268 if (!tmp1 || rb_is_black(tmp1)) { 269 tmp2 = sibling->rb_left; 270 if (!tmp2 || rb_is_black(tmp2)) { 271 /* 272 * Case 2 - sibling color flip 273 * (p could be either color here) 274 * 275 * (p) (p) 276 * / \ / \ 277 * N S --> N s 278 * / \ / \ 279 * Sl Sr Sl Sr 280 * 281 * This leaves us violating 5) which 282 * can be fixed by flipping p to black 283 * if it was red, or by recursing at p. 284 * p is red when coming from Case 1. 285 */ 286 rb_set_parent_color(sibling, parent, 287 RB_RED); 288 if (rb_is_red(parent)) 289 rb_set_black(parent); 290 else { 291 node = parent; 292 parent = rb_parent(node); 293 if (parent) 294 continue; 295 } 296 break; 297 } 298 /* 299 * Case 3 - right rotate at sibling 300 * (p could be either color here) 301 * 302 * (p) (p) 303 * / \ / \ 304 * N S --> N Sl 305 * / \ \ 306 * sl Sr s 307 * \ 308 * Sr 309 */ 310 sibling->rb_left = tmp1 = tmp2->rb_right; 311 tmp2->rb_right = sibling; 312 parent->rb_right = tmp2; 313 if (tmp1) 314 rb_set_parent_color(tmp1, sibling, 315 RB_BLACK); 316 augment->rotate(sibling, tmp2); 317 tmp1 = sibling; 318 sibling = tmp2; 319 } 320 /* 321 * Case 4 - left rotate at parent + color flips 322 * (p and sl could be either color here. 323 * After rotation, p becomes black, s acquires 324 * p's color, and sl keeps its color) 325 * 326 * (p) (s) 327 * / \ / \ 328 * N S --> P Sr 329 * / \ / \ 330 * (sl) sr N (sl) 331 */ 332 parent->rb_right = tmp2 = sibling->rb_left; 333 sibling->rb_left = parent; 334 rb_set_parent_color(tmp1, sibling, RB_BLACK); 335 if (tmp2) 336 rb_set_parent(tmp2, parent); 337 __rb_rotate_set_parents(parent, sibling, root, 338 RB_BLACK); 339 augment->rotate(parent, sibling); 340 break; 341 } else { 342 sibling = parent->rb_left; 343 if (rb_is_red(sibling)) { 344 /* Case 1 - right rotate at parent */ 345 parent->rb_left = tmp1 = sibling->rb_right; 346 sibling->rb_right = parent; 347 rb_set_parent_color(tmp1, parent, RB_BLACK); 348 __rb_rotate_set_parents(parent, sibling, root, 349 RB_RED); 350 augment->rotate(parent, sibling); 351 sibling = tmp1; 352 } 353 tmp1 = sibling->rb_left; 354 if (!tmp1 || rb_is_black(tmp1)) { 355 tmp2 = sibling->rb_right; 356 if (!tmp2 || rb_is_black(tmp2)) { 357 /* Case 2 - sibling color flip */ 358 rb_set_parent_color(sibling, parent, 359 RB_RED); 360 if (rb_is_red(parent)) 361 rb_set_black(parent); 362 else { 363 node = parent; 364 parent = rb_parent(node); 365 if (parent) 366 continue; 367 } 368 break; 369 } 370 /* Case 3 - right rotate at sibling */ 371 sibling->rb_right = tmp1 = tmp2->rb_left; 372 tmp2->rb_left = sibling; 373 parent->rb_left = tmp2; 374 if (tmp1) 375 rb_set_parent_color(tmp1, sibling, 376 RB_BLACK); 377 augment->rotate(sibling, tmp2); 378 tmp1 = sibling; 379 sibling = tmp2; 380 } 381 /* Case 4 - left rotate at parent + color flips */ 382 parent->rb_left = tmp2 = sibling->rb_right; 383 sibling->rb_right = parent; 384 rb_set_parent_color(tmp1, sibling, RB_BLACK); 385 if (tmp2) 386 rb_set_parent(tmp2, parent); 387 __rb_rotate_set_parents(parent, sibling, root, 388 RB_BLACK); 389 augment->rotate(parent, sibling); 390 break; 391 } 392 } 393 } 394 395 static __always_inline void 396 __rb_erase(struct rb_node *node, struct rb_root *root, 397 const struct rb_augment_callbacks *augment) 398 { 399 struct rb_node *child = node->rb_right, *tmp = node->rb_left; 400 struct rb_node *parent, *rebalance; 401 unsigned long pc; 402 403 if (!tmp) { 404 /* 405 * Case 1: node to erase has no more than 1 child (easy!) 406 * 407 * Note that if there is one child it must be red due to 5) 408 * and node must be black due to 4). We adjust colors locally 409 * so as to bypass __rb_erase_color() later on. 410 */ 411 pc = node->__rb_parent_color; 412 parent = __rb_parent(pc); 413 __rb_change_child(node, child, parent, root); 414 if (child) { 415 child->__rb_parent_color = pc; 416 rebalance = NULL; 417 } else 418 rebalance = __rb_is_black(pc) ? parent : NULL; 419 tmp = parent; 420 } else if (!child) { 421 /* Still case 1, but this time the child is node->rb_left */ 422 tmp->__rb_parent_color = pc = node->__rb_parent_color; 423 parent = __rb_parent(pc); 424 __rb_change_child(node, tmp, parent, root); 425 rebalance = NULL; 426 tmp = parent; 427 } else { 428 struct rb_node *successor = child, *child2; 429 tmp = child->rb_left; 430 if (!tmp) { 431 /* 432 * Case 2: node's successor is its right child 433 * 434 * (n) (s) 435 * / \ / \ 436 * (x) (s) -> (x) (c) 437 * \ 438 * (c) 439 */ 440 parent = successor; 441 child2 = successor->rb_right; 442 augment->copy(node, successor); 443 } else { 444 /* 445 * Case 3: node's successor is leftmost under 446 * node's right child subtree 447 * 448 * (n) (s) 449 * / \ / \ 450 * (x) (y) -> (x) (y) 451 * / / 452 * (p) (p) 453 * / / 454 * (s) (c) 455 * \ 456 * (c) 457 */ 458 do { 459 parent = successor; 460 successor = tmp; 461 tmp = tmp->rb_left; 462 } while (tmp); 463 parent->rb_left = child2 = successor->rb_right; 464 successor->rb_right = child; 465 rb_set_parent(child, successor); 466 augment->copy(node, successor); 467 augment->propagate(parent, successor); 468 } 469 470 successor->rb_left = tmp = node->rb_left; 471 rb_set_parent(tmp, successor); 472 473 pc = node->__rb_parent_color; 474 tmp = __rb_parent(pc); 475 __rb_change_child(node, successor, tmp, root); 476 if (child2) { 477 successor->__rb_parent_color = pc; 478 rb_set_parent_color(child2, parent, RB_BLACK); 479 rebalance = NULL; 480 } else { 481 unsigned long pc2 = successor->__rb_parent_color; 482 successor->__rb_parent_color = pc; 483 rebalance = __rb_is_black(pc2) ? parent : NULL; 484 } 485 tmp = successor; 486 } 487 488 augment->propagate(tmp, NULL); 489 if (rebalance) 490 __rb_erase_color(rebalance, root, augment); 491 } 492 493 /* 494 * Non-augmented rbtree manipulation functions. 495 * 496 * We use dummy augmented callbacks here, and have the compiler optimize them 497 * out of the rb_insert_color() and rb_erase() function definitions. 498 */ 499 500 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {} 501 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {} 502 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {} 503 504 static const struct rb_augment_callbacks dummy_callbacks = { 505 dummy_propagate, dummy_copy, dummy_rotate 506 }; 507 508 void rb_insert_color(struct rb_node *node, struct rb_root *root) 509 { 510 __rb_insert(node, root, dummy_rotate); 511 } 512 EXPORT_SYMBOL(rb_insert_color); 513 514 void rb_erase(struct rb_node *node, struct rb_root *root) 515 { 516 __rb_erase(node, root, &dummy_callbacks); 517 } 518 EXPORT_SYMBOL(rb_erase); 519 520 /* 521 * Augmented rbtree manipulation functions. 522 * 523 * This instantiates the same __always_inline functions as in the non-augmented 524 * case, but this time with user-defined callbacks. 525 */ 526 527 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, 528 void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 529 { 530 __rb_insert(node, root, augment_rotate); 531 } 532 EXPORT_SYMBOL(__rb_insert_augmented); 533 534 void rb_erase_augmented(struct rb_node *node, struct rb_root *root, 535 const struct rb_augment_callbacks *augment) 536 { 537 __rb_erase(node, root, augment); 538 } 539 EXPORT_SYMBOL(rb_erase_augmented); 540 541 /* 542 * This function returns the first node (in sort order) of the tree. 543 */ 544 struct rb_node *rb_first(const struct rb_root *root) 545 { 546 struct rb_node *n; 547 548 n = root->rb_node; 549 if (!n) 550 return NULL; 551 while (n->rb_left) 552 n = n->rb_left; 553 return n; 554 } 555 EXPORT_SYMBOL(rb_first); 556 557 struct rb_node *rb_last(const struct rb_root *root) 558 { 559 struct rb_node *n; 560 561 n = root->rb_node; 562 if (!n) 563 return NULL; 564 while (n->rb_right) 565 n = n->rb_right; 566 return n; 567 } 568 EXPORT_SYMBOL(rb_last); 569 570 struct rb_node *rb_next(const struct rb_node *node) 571 { 572 struct rb_node *parent; 573 574 if (RB_EMPTY_NODE(node)) 575 return NULL; 576 577 /* 578 * If we have a right-hand child, go down and then left as far 579 * as we can. 580 */ 581 if (node->rb_right) { 582 node = node->rb_right; 583 while (node->rb_left) 584 node=node->rb_left; 585 return (struct rb_node *)node; 586 } 587 588 /* 589 * No right-hand children. Everything down and left is smaller than us, 590 * so any 'next' node must be in the general direction of our parent. 591 * Go up the tree; any time the ancestor is a right-hand child of its 592 * parent, keep going up. First time it's a left-hand child of its 593 * parent, said parent is our 'next' node. 594 */ 595 while ((parent = rb_parent(node)) && node == parent->rb_right) 596 node = parent; 597 598 return parent; 599 } 600 EXPORT_SYMBOL(rb_next); 601 602 struct rb_node *rb_prev(const struct rb_node *node) 603 { 604 struct rb_node *parent; 605 606 if (RB_EMPTY_NODE(node)) 607 return NULL; 608 609 /* 610 * If we have a left-hand child, go down and then right as far 611 * as we can. 612 */ 613 if (node->rb_left) { 614 node = node->rb_left; 615 while (node->rb_right) 616 node=node->rb_right; 617 return (struct rb_node *)node; 618 } 619 620 /* 621 * No left-hand children. Go up till we find an ancestor which 622 * is a right-hand child of its parent. 623 */ 624 while ((parent = rb_parent(node)) && node == parent->rb_left) 625 node = parent; 626 627 return parent; 628 } 629 EXPORT_SYMBOL(rb_prev); 630 631 void rb_replace_node(struct rb_node *victim, struct rb_node *new, 632 struct rb_root *root) 633 { 634 struct rb_node *parent = rb_parent(victim); 635 636 /* Set the surrounding nodes to point to the replacement */ 637 __rb_change_child(victim, new, parent, root); 638 if (victim->rb_left) 639 rb_set_parent(victim->rb_left, new); 640 if (victim->rb_right) 641 rb_set_parent(victim->rb_right, new); 642 643 /* Copy the pointers/colour from the victim to the replacement */ 644 *new = *victim; 645 } 646 EXPORT_SYMBOL(rb_replace_node); 647