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 __always_inline void 198 __rb_erase_color(struct rb_node *parent, struct rb_root *root, 199 void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 200 { 201 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; 202 203 while (true) { 204 /* 205 * Loop invariants: 206 * - node is black (or NULL on first iteration) 207 * - node is not the root (parent is not NULL) 208 * - All leaf paths going through parent and node have a 209 * black node count that is 1 lower than other leaf paths. 210 */ 211 sibling = parent->rb_right; 212 if (node != sibling) { /* node == parent->rb_left */ 213 if (rb_is_red(sibling)) { 214 /* 215 * Case 1 - left rotate at parent 216 * 217 * P S 218 * / \ / \ 219 * N s --> p Sr 220 * / \ / \ 221 * Sl Sr N Sl 222 */ 223 parent->rb_right = tmp1 = sibling->rb_left; 224 sibling->rb_left = parent; 225 rb_set_parent_color(tmp1, parent, RB_BLACK); 226 __rb_rotate_set_parents(parent, sibling, root, 227 RB_RED); 228 augment_rotate(parent, sibling); 229 sibling = tmp1; 230 } 231 tmp1 = sibling->rb_right; 232 if (!tmp1 || rb_is_black(tmp1)) { 233 tmp2 = sibling->rb_left; 234 if (!tmp2 || rb_is_black(tmp2)) { 235 /* 236 * Case 2 - sibling color flip 237 * (p could be either color here) 238 * 239 * (p) (p) 240 * / \ / \ 241 * N S --> N s 242 * / \ / \ 243 * Sl Sr Sl Sr 244 * 245 * This leaves us violating 5) which 246 * can be fixed by flipping p to black 247 * if it was red, or by recursing at p. 248 * p is red when coming from Case 1. 249 */ 250 rb_set_parent_color(sibling, parent, 251 RB_RED); 252 if (rb_is_red(parent)) 253 rb_set_black(parent); 254 else { 255 node = parent; 256 parent = rb_parent(node); 257 if (parent) 258 continue; 259 } 260 break; 261 } 262 /* 263 * Case 3 - right rotate at sibling 264 * (p could be either color here) 265 * 266 * (p) (p) 267 * / \ / \ 268 * N S --> N Sl 269 * / \ \ 270 * sl Sr s 271 * \ 272 * Sr 273 */ 274 sibling->rb_left = tmp1 = tmp2->rb_right; 275 tmp2->rb_right = sibling; 276 parent->rb_right = tmp2; 277 if (tmp1) 278 rb_set_parent_color(tmp1, sibling, 279 RB_BLACK); 280 augment_rotate(sibling, tmp2); 281 tmp1 = sibling; 282 sibling = tmp2; 283 } 284 /* 285 * Case 4 - left rotate at parent + color flips 286 * (p and sl could be either color here. 287 * After rotation, p becomes black, s acquires 288 * p's color, and sl keeps its color) 289 * 290 * (p) (s) 291 * / \ / \ 292 * N S --> P Sr 293 * / \ / \ 294 * (sl) sr N (sl) 295 */ 296 parent->rb_right = tmp2 = sibling->rb_left; 297 sibling->rb_left = parent; 298 rb_set_parent_color(tmp1, sibling, RB_BLACK); 299 if (tmp2) 300 rb_set_parent(tmp2, parent); 301 __rb_rotate_set_parents(parent, sibling, root, 302 RB_BLACK); 303 augment_rotate(parent, sibling); 304 break; 305 } else { 306 sibling = parent->rb_left; 307 if (rb_is_red(sibling)) { 308 /* Case 1 - right rotate at parent */ 309 parent->rb_left = tmp1 = sibling->rb_right; 310 sibling->rb_right = parent; 311 rb_set_parent_color(tmp1, parent, RB_BLACK); 312 __rb_rotate_set_parents(parent, sibling, root, 313 RB_RED); 314 augment_rotate(parent, sibling); 315 sibling = tmp1; 316 } 317 tmp1 = sibling->rb_left; 318 if (!tmp1 || rb_is_black(tmp1)) { 319 tmp2 = sibling->rb_right; 320 if (!tmp2 || rb_is_black(tmp2)) { 321 /* Case 2 - sibling color flip */ 322 rb_set_parent_color(sibling, parent, 323 RB_RED); 324 if (rb_is_red(parent)) 325 rb_set_black(parent); 326 else { 327 node = parent; 328 parent = rb_parent(node); 329 if (parent) 330 continue; 331 } 332 break; 333 } 334 /* Case 3 - right rotate at sibling */ 335 sibling->rb_right = tmp1 = tmp2->rb_left; 336 tmp2->rb_left = sibling; 337 parent->rb_left = tmp2; 338 if (tmp1) 339 rb_set_parent_color(tmp1, sibling, 340 RB_BLACK); 341 augment_rotate(sibling, tmp2); 342 tmp1 = sibling; 343 sibling = tmp2; 344 } 345 /* Case 4 - left rotate at parent + color flips */ 346 parent->rb_left = tmp2 = sibling->rb_right; 347 sibling->rb_right = parent; 348 rb_set_parent_color(tmp1, sibling, RB_BLACK); 349 if (tmp2) 350 rb_set_parent(tmp2, parent); 351 __rb_rotate_set_parents(parent, sibling, root, 352 RB_BLACK); 353 augment_rotate(parent, sibling); 354 break; 355 } 356 } 357 } 358 EXPORT_SYMBOL(__rb_erase_color); 359 360 /* 361 * Non-augmented rbtree manipulation functions. 362 * 363 * We use dummy augmented callbacks here, and have the compiler optimize them 364 * out of the rb_insert_color() and rb_erase() function definitions. 365 */ 366 367 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {} 368 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {} 369 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {} 370 371 static const struct rb_augment_callbacks dummy_callbacks = { 372 dummy_propagate, dummy_copy, dummy_rotate 373 }; 374 375 void rb_insert_color(struct rb_node *node, struct rb_root *root) 376 { 377 __rb_insert(node, root, dummy_rotate); 378 } 379 EXPORT_SYMBOL(rb_insert_color); 380 381 void rb_erase(struct rb_node *node, struct rb_root *root) 382 { 383 rb_erase_augmented(node, root, &dummy_callbacks); 384 } 385 EXPORT_SYMBOL(rb_erase); 386 387 /* 388 * Augmented rbtree manipulation functions. 389 * 390 * This instantiates the same __always_inline functions as in the non-augmented 391 * case, but this time with user-defined callbacks. 392 */ 393 394 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, 395 void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 396 { 397 __rb_insert(node, root, augment_rotate); 398 } 399 EXPORT_SYMBOL(__rb_insert_augmented); 400 401 /* 402 * This function returns the first node (in sort order) of the tree. 403 */ 404 struct rb_node *rb_first(const struct rb_root *root) 405 { 406 struct rb_node *n; 407 408 n = root->rb_node; 409 if (!n) 410 return NULL; 411 while (n->rb_left) 412 n = n->rb_left; 413 return n; 414 } 415 EXPORT_SYMBOL(rb_first); 416 417 struct rb_node *rb_last(const struct rb_root *root) 418 { 419 struct rb_node *n; 420 421 n = root->rb_node; 422 if (!n) 423 return NULL; 424 while (n->rb_right) 425 n = n->rb_right; 426 return n; 427 } 428 EXPORT_SYMBOL(rb_last); 429 430 struct rb_node *rb_next(const struct rb_node *node) 431 { 432 struct rb_node *parent; 433 434 if (RB_EMPTY_NODE(node)) 435 return NULL; 436 437 /* 438 * If we have a right-hand child, go down and then left as far 439 * as we can. 440 */ 441 if (node->rb_right) { 442 node = node->rb_right; 443 while (node->rb_left) 444 node=node->rb_left; 445 return (struct rb_node *)node; 446 } 447 448 /* 449 * No right-hand children. Everything down and left is smaller than us, 450 * so any 'next' node must be in the general direction of our parent. 451 * Go up the tree; any time the ancestor is a right-hand child of its 452 * parent, keep going up. First time it's a left-hand child of its 453 * parent, said parent is our 'next' node. 454 */ 455 while ((parent = rb_parent(node)) && node == parent->rb_right) 456 node = parent; 457 458 return parent; 459 } 460 EXPORT_SYMBOL(rb_next); 461 462 struct rb_node *rb_prev(const struct rb_node *node) 463 { 464 struct rb_node *parent; 465 466 if (RB_EMPTY_NODE(node)) 467 return NULL; 468 469 /* 470 * If we have a left-hand child, go down and then right as far 471 * as we can. 472 */ 473 if (node->rb_left) { 474 node = node->rb_left; 475 while (node->rb_right) 476 node=node->rb_right; 477 return (struct rb_node *)node; 478 } 479 480 /* 481 * No left-hand children. Go up till we find an ancestor which 482 * is a right-hand child of its parent. 483 */ 484 while ((parent = rb_parent(node)) && node == parent->rb_left) 485 node = parent; 486 487 return parent; 488 } 489 EXPORT_SYMBOL(rb_prev); 490 491 void rb_replace_node(struct rb_node *victim, struct rb_node *new, 492 struct rb_root *root) 493 { 494 struct rb_node *parent = rb_parent(victim); 495 496 /* Set the surrounding nodes to point to the replacement */ 497 __rb_change_child(victim, new, parent, root); 498 if (victim->rb_left) 499 rb_set_parent(victim->rb_left, new); 500 if (victim->rb_right) 501 rb_set_parent(victim->rb_right, new); 502 503 /* Copy the pointers/colour from the victim to the replacement */ 504 *new = *victim; 505 } 506 EXPORT_SYMBOL(rb_replace_node); 507