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 rb->__rb_parent_color |= RB_BLACK; 50 } 51 52 static inline struct rb_node *rb_red_parent(struct rb_node *red) 53 { 54 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 struct rb_node *parent = rb_parent(old); 67 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 __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 if (!parent) { 87 rb_set_parent_color(node, NULL, RB_BLACK); 88 break; 89 } else if (rb_is_black(parent)) 90 break; 91 92 gparent = rb_red_parent(parent); 93 94 tmp = gparent->rb_right; 95 if (parent != tmp) { /* parent == gparent->rb_left */ 96 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 parent = rb_parent(node); 114 rb_set_parent_color(node, parent, RB_RED); 115 continue; 116 } 117 118 tmp = parent->rb_right; 119 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 parent->rb_right = tmp = node->rb_left; 133 node->rb_left = parent; 134 if (tmp) 135 rb_set_parent_color(tmp, parent, 136 RB_BLACK); 137 rb_set_parent_color(parent, node, RB_RED); 138 augment_rotate(parent, node); 139 parent = node; 140 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 gparent->rb_left = tmp; /* == parent->rb_right */ 153 parent->rb_right = gparent; 154 if (tmp) 155 rb_set_parent_color(tmp, gparent, RB_BLACK); 156 __rb_rotate_set_parents(gparent, parent, root, RB_RED); 157 augment_rotate(gparent, parent); 158 break; 159 } else { 160 tmp = gparent->rb_left; 161 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 parent = rb_parent(node); 167 rb_set_parent_color(node, parent, RB_RED); 168 continue; 169 } 170 171 tmp = parent->rb_left; 172 if (node == tmp) { 173 /* Case 2 - right rotate at parent */ 174 parent->rb_left = tmp = node->rb_right; 175 node->rb_right = parent; 176 if (tmp) 177 rb_set_parent_color(tmp, parent, 178 RB_BLACK); 179 rb_set_parent_color(parent, node, RB_RED); 180 augment_rotate(parent, node); 181 parent = node; 182 tmp = node->rb_left; 183 } 184 185 /* Case 3 - left rotate at gparent */ 186 gparent->rb_right = tmp; /* == parent->rb_left */ 187 parent->rb_left = gparent; 188 if (tmp) 189 rb_set_parent_color(tmp, gparent, RB_BLACK); 190 __rb_rotate_set_parents(gparent, parent, root, RB_RED); 191 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 sibling = parent->rb_right; 216 if (node != sibling) { /* node == parent->rb_left */ 217 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 parent->rb_right = tmp1 = sibling->rb_left; 228 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 augment_rotate(parent, sibling); 233 sibling = tmp1; 234 } 235 tmp1 = sibling->rb_right; 236 if (!tmp1 || rb_is_black(tmp1)) { 237 tmp2 = sibling->rb_left; 238 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 if (rb_is_red(parent)) 257 rb_set_black(parent); 258 else { 259 node = parent; 260 parent = rb_parent(node); 261 if (parent) 262 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 sibling->rb_left = tmp1 = tmp2->rb_right; 279 tmp2->rb_right = sibling; 280 parent->rb_right = tmp2; 281 if (tmp1) 282 rb_set_parent_color(tmp1, sibling, 283 RB_BLACK); 284 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 parent->rb_right = tmp2 = sibling->rb_left; 301 sibling->rb_left = parent; 302 rb_set_parent_color(tmp1, sibling, RB_BLACK); 303 if (tmp2) 304 rb_set_parent(tmp2, parent); 305 __rb_rotate_set_parents(parent, sibling, root, 306 RB_BLACK); 307 augment_rotate(parent, sibling); 308 break; 309 } else { 310 sibling = parent->rb_left; 311 if (rb_is_red(sibling)) { 312 /* Case 1 - right rotate at parent */ 313 parent->rb_left = tmp1 = sibling->rb_right; 314 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 augment_rotate(parent, sibling); 319 sibling = tmp1; 320 } 321 tmp1 = sibling->rb_left; 322 if (!tmp1 || rb_is_black(tmp1)) { 323 tmp2 = sibling->rb_right; 324 if (!tmp2 || rb_is_black(tmp2)) { 325 /* Case 2 - sibling color flip */ 326 rb_set_parent_color(sibling, parent, 327 RB_RED); 328 if (rb_is_red(parent)) 329 rb_set_black(parent); 330 else { 331 node = parent; 332 parent = rb_parent(node); 333 if (parent) 334 continue; 335 } 336 break; 337 } 338 /* Case 3 - right rotate at sibling */ 339 sibling->rb_right = tmp1 = tmp2->rb_left; 340 tmp2->rb_left = sibling; 341 parent->rb_left = tmp2; 342 if (tmp1) 343 rb_set_parent_color(tmp1, sibling, 344 RB_BLACK); 345 augment_rotate(sibling, tmp2); 346 tmp1 = sibling; 347 sibling = tmp2; 348 } 349 /* Case 4 - left rotate at parent + color flips */ 350 parent->rb_left = tmp2 = sibling->rb_right; 351 sibling->rb_right = parent; 352 rb_set_parent_color(tmp1, sibling, RB_BLACK); 353 if (tmp2) 354 rb_set_parent(tmp2, parent); 355 __rb_rotate_set_parents(parent, sibling, root, 356 RB_BLACK); 357 augment_rotate(parent, sibling); 358 break; 359 } 360 } 361 } 362 363 /* Non-inline version for rb_erase_augmented() use */ 364 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 } 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 void rb_insert_color(struct rb_node *node, struct rb_root *root) 387 { 388 __rb_insert(node, root, dummy_rotate); 389 } 390 EXPORT_SYMBOL(rb_insert_color); 391 392 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 if (rebalance) 397 ____rb_erase_color(rebalance, root, dummy_rotate); 398 } 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 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 } 413 EXPORT_SYMBOL(__rb_insert_augmented); 414 415 /* 416 * This function returns the first node (in sort order) of the tree. 417 */ 418 struct rb_node *rb_first(const struct rb_root *root) 419 { 420 struct rb_node *n; 421 422 n = root->rb_node; 423 if (!n) 424 return NULL; 425 while (n->rb_left) 426 n = n->rb_left; 427 return n; 428 } 429 EXPORT_SYMBOL(rb_first); 430 431 struct rb_node *rb_last(const struct rb_root *root) 432 { 433 struct rb_node *n; 434 435 n = root->rb_node; 436 if (!n) 437 return NULL; 438 while (n->rb_right) 439 n = n->rb_right; 440 return n; 441 } 442 EXPORT_SYMBOL(rb_last); 443 444 struct rb_node *rb_next(const struct rb_node *node) 445 { 446 struct rb_node *parent; 447 448 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 if (node->rb_right) { 456 node = node->rb_right; 457 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 while ((parent = rb_parent(node)) && node == parent->rb_right) 470 node = parent; 471 472 return parent; 473 } 474 EXPORT_SYMBOL(rb_next); 475 476 struct rb_node *rb_prev(const struct rb_node *node) 477 { 478 struct rb_node *parent; 479 480 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 if (node->rb_left) { 488 node = node->rb_left; 489 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 while ((parent = rb_parent(node)) && node == parent->rb_left) 499 node = parent; 500 501 return parent; 502 } 503 EXPORT_SYMBOL(rb_prev); 504 505 void rb_replace_node(struct rb_node *victim, struct rb_node *new, 506 struct rb_root *root) 507 { 508 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 if (victim->rb_left) 513 rb_set_parent(victim->rb_left, new); 514 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 *new = *victim; 519 } 520 EXPORT_SYMBOL(rb_replace_node); 521