1 /*------------------------------------------------------------------------- 2 * 3 * rbtree.c 4 * implementation for PostgreSQL generic Red-Black binary tree package 5 * Adopted from http://algolist.manual.ru/ds/rbtree.php 6 * 7 * This code comes from Thomas Niemann's "Sorting and Searching Algorithms: 8 * a Cookbook". 9 * 10 * See http://www.cs.auckland.ac.nz/software/AlgAnim/niemann/s_man.htm for 11 * license terms: "Source code, when part of a software project, may be used 12 * freely without reference to the author." 13 * 14 * Red-black trees are a type of balanced binary tree wherein (1) any child of 15 * a red node is always black, and (2) every path from root to leaf traverses 16 * an equal number of black nodes. From these properties, it follows that the 17 * longest path from root to leaf is only about twice as long as the shortest, 18 * so lookups are guaranteed to run in O(lg n) time. 19 * 20 * Copyright (c) 2009-2021, PostgreSQL Global Development Group 21 * 22 * IDENTIFICATION 23 * src/backend/lib/rbtree.c 24 * 25 *------------------------------------------------------------------------- 26 */ 27 #include "postgres.h" 28 29 #include "lib/rbtree.h" 30 31 32 /* 33 * Colors of nodes (values of RBTNode.color) 34 */ 35 #define RBTBLACK (0) 36 #define RBTRED (1) 37 38 /* 39 * RBTree control structure 40 */ 41 struct RBTree 42 { 43 RBTNode *root; /* root node, or RBTNIL if tree is empty */ 44 45 /* Remaining fields are constant after rbt_create */ 46 47 Size node_size; /* actual size of tree nodes */ 48 /* The caller-supplied manipulation functions */ 49 rbt_comparator comparator; 50 rbt_combiner combiner; 51 rbt_allocfunc allocfunc; 52 rbt_freefunc freefunc; 53 /* Passthrough arg passed to all manipulation functions */ 54 void *arg; 55 }; 56 57 /* 58 * all leafs are sentinels, use customized NIL name to prevent 59 * collision with system-wide constant NIL which is actually NULL 60 */ 61 #define RBTNIL (&sentinel) 62 63 static RBTNode sentinel = 64 { 65 RBTBLACK, RBTNIL, RBTNIL, NULL 66 }; 67 68 69 /* 70 * rbt_create: create an empty RBTree 71 * 72 * Arguments are: 73 * node_size: actual size of tree nodes (> sizeof(RBTNode)) 74 * The manipulation functions: 75 * comparator: compare two RBTNodes for less/equal/greater 76 * combiner: merge an existing tree entry with a new one 77 * allocfunc: allocate a new RBTNode 78 * freefunc: free an old RBTNode 79 * arg: passthrough pointer that will be passed to the manipulation functions 80 * 81 * Note that the combiner's righthand argument will be a "proposed" tree node, 82 * ie the input to rbt_insert, in which the RBTNode fields themselves aren't 83 * valid. Similarly, either input to the comparator may be a "proposed" node. 84 * This shouldn't matter since the functions aren't supposed to look at the 85 * RBTNode fields, only the extra fields of the struct the RBTNode is embedded 86 * in. 87 * 88 * The freefunc should just be pfree or equivalent; it should NOT attempt 89 * to free any subsidiary data, because the node passed to it may not contain 90 * valid data! freefunc can be NULL if caller doesn't require retail 91 * space reclamation. 92 * 93 * The RBTree node is palloc'd in the caller's memory context. Note that 94 * all contents of the tree are actually allocated by the caller, not here. 95 * 96 * Since tree contents are managed by the caller, there is currently not 97 * an explicit "destroy" operation; typically a tree would be freed by 98 * resetting or deleting the memory context it's stored in. You can pfree 99 * the RBTree node if you feel the urge. 100 */ 101 RBTree * 102 rbt_create(Size node_size, 103 rbt_comparator comparator, 104 rbt_combiner combiner, 105 rbt_allocfunc allocfunc, 106 rbt_freefunc freefunc, 107 void *arg) 108 { 109 RBTree *tree = (RBTree *) palloc(sizeof(RBTree)); 110 111 Assert(node_size > sizeof(RBTNode)); 112 113 tree->root = RBTNIL; 114 tree->node_size = node_size; 115 tree->comparator = comparator; 116 tree->combiner = combiner; 117 tree->allocfunc = allocfunc; 118 tree->freefunc = freefunc; 119 120 tree->arg = arg; 121 122 return tree; 123 } 124 125 /* Copy the additional data fields from one RBTNode to another */ 126 static inline void 127 rbt_copy_data(RBTree *rbt, RBTNode *dest, const RBTNode *src) 128 { 129 memcpy(dest + 1, src + 1, rbt->node_size - sizeof(RBTNode)); 130 } 131 132 /********************************************************************** 133 * Search * 134 **********************************************************************/ 135 136 /* 137 * rbt_find: search for a value in an RBTree 138 * 139 * data represents the value to try to find. Its RBTNode fields need not 140 * be valid, it's the extra data in the larger struct that is of interest. 141 * 142 * Returns the matching tree entry, or NULL if no match is found. 143 */ 144 RBTNode * 145 rbt_find(RBTree *rbt, const RBTNode *data) 146 { 147 RBTNode *node = rbt->root; 148 149 while (node != RBTNIL) 150 { 151 int cmp = rbt->comparator(data, node, rbt->arg); 152 153 if (cmp == 0) 154 return node; 155 else if (cmp < 0) 156 node = node->left; 157 else 158 node = node->right; 159 } 160 161 return NULL; 162 } 163 164 /* 165 * rbt_leftmost: fetch the leftmost (smallest-valued) tree node. 166 * Returns NULL if tree is empty. 167 * 168 * Note: in the original implementation this included an unlink step, but 169 * that's a bit awkward. Just call rbt_delete on the result if that's what 170 * you want. 171 */ 172 RBTNode * 173 rbt_leftmost(RBTree *rbt) 174 { 175 RBTNode *node = rbt->root; 176 RBTNode *leftmost = rbt->root; 177 178 while (node != RBTNIL) 179 { 180 leftmost = node; 181 node = node->left; 182 } 183 184 if (leftmost != RBTNIL) 185 return leftmost; 186 187 return NULL; 188 } 189 190 /********************************************************************** 191 * Insertion * 192 **********************************************************************/ 193 194 /* 195 * Rotate node x to left. 196 * 197 * x's right child takes its place in the tree, and x becomes the left 198 * child of that node. 199 */ 200 static void 201 rbt_rotate_left(RBTree *rbt, RBTNode *x) 202 { 203 RBTNode *y = x->right; 204 205 /* establish x->right link */ 206 x->right = y->left; 207 if (y->left != RBTNIL) 208 y->left->parent = x; 209 210 /* establish y->parent link */ 211 if (y != RBTNIL) 212 y->parent = x->parent; 213 if (x->parent) 214 { 215 if (x == x->parent->left) 216 x->parent->left = y; 217 else 218 x->parent->right = y; 219 } 220 else 221 { 222 rbt->root = y; 223 } 224 225 /* link x and y */ 226 y->left = x; 227 if (x != RBTNIL) 228 x->parent = y; 229 } 230 231 /* 232 * Rotate node x to right. 233 * 234 * x's left right child takes its place in the tree, and x becomes the right 235 * child of that node. 236 */ 237 static void 238 rbt_rotate_right(RBTree *rbt, RBTNode *x) 239 { 240 RBTNode *y = x->left; 241 242 /* establish x->left link */ 243 x->left = y->right; 244 if (y->right != RBTNIL) 245 y->right->parent = x; 246 247 /* establish y->parent link */ 248 if (y != RBTNIL) 249 y->parent = x->parent; 250 if (x->parent) 251 { 252 if (x == x->parent->right) 253 x->parent->right = y; 254 else 255 x->parent->left = y; 256 } 257 else 258 { 259 rbt->root = y; 260 } 261 262 /* link x and y */ 263 y->right = x; 264 if (x != RBTNIL) 265 x->parent = y; 266 } 267 268 /* 269 * Maintain Red-Black tree balance after inserting node x. 270 * 271 * The newly inserted node is always initially marked red. That may lead to 272 * a situation where a red node has a red child, which is prohibited. We can 273 * always fix the problem by a series of color changes and/or "rotations", 274 * which move the problem progressively higher up in the tree. If one of the 275 * two red nodes is the root, we can always fix the problem by changing the 276 * root from red to black. 277 * 278 * (This does not work lower down in the tree because we must also maintain 279 * the invariant that every leaf has equal black-height.) 280 */ 281 static void 282 rbt_insert_fixup(RBTree *rbt, RBTNode *x) 283 { 284 /* 285 * x is always a red node. Initially, it is the newly inserted node. Each 286 * iteration of this loop moves it higher up in the tree. 287 */ 288 while (x != rbt->root && x->parent->color == RBTRED) 289 { 290 /* 291 * x and x->parent are both red. Fix depends on whether x->parent is 292 * a left or right child. In either case, we define y to be the 293 * "uncle" of x, that is, the other child of x's grandparent. 294 * 295 * If the uncle is red, we flip the grandparent to red and its two 296 * children to black. Then we loop around again to check whether the 297 * grandparent still has a problem. 298 * 299 * If the uncle is black, we will perform one or two "rotations" to 300 * balance the tree. Either x or x->parent will take the 301 * grandparent's position in the tree and recolored black, and the 302 * original grandparent will be recolored red and become a child of 303 * that node. This always leaves us with a valid red-black tree, so 304 * the loop will terminate. 305 */ 306 if (x->parent == x->parent->parent->left) 307 { 308 RBTNode *y = x->parent->parent->right; 309 310 if (y->color == RBTRED) 311 { 312 /* uncle is RBTRED */ 313 x->parent->color = RBTBLACK; 314 y->color = RBTBLACK; 315 x->parent->parent->color = RBTRED; 316 317 x = x->parent->parent; 318 } 319 else 320 { 321 /* uncle is RBTBLACK */ 322 if (x == x->parent->right) 323 { 324 /* make x a left child */ 325 x = x->parent; 326 rbt_rotate_left(rbt, x); 327 } 328 329 /* recolor and rotate */ 330 x->parent->color = RBTBLACK; 331 x->parent->parent->color = RBTRED; 332 333 rbt_rotate_right(rbt, x->parent->parent); 334 } 335 } 336 else 337 { 338 /* mirror image of above code */ 339 RBTNode *y = x->parent->parent->left; 340 341 if (y->color == RBTRED) 342 { 343 /* uncle is RBTRED */ 344 x->parent->color = RBTBLACK; 345 y->color = RBTBLACK; 346 x->parent->parent->color = RBTRED; 347 348 x = x->parent->parent; 349 } 350 else 351 { 352 /* uncle is RBTBLACK */ 353 if (x == x->parent->left) 354 { 355 x = x->parent; 356 rbt_rotate_right(rbt, x); 357 } 358 x->parent->color = RBTBLACK; 359 x->parent->parent->color = RBTRED; 360 361 rbt_rotate_left(rbt, x->parent->parent); 362 } 363 } 364 } 365 366 /* 367 * The root may already have been black; if not, the black-height of every 368 * node in the tree increases by one. 369 */ 370 rbt->root->color = RBTBLACK; 371 } 372 373 /* 374 * rbt_insert: insert a new value into the tree. 375 * 376 * data represents the value to insert. Its RBTNode fields need not 377 * be valid, it's the extra data in the larger struct that is of interest. 378 * 379 * If the value represented by "data" is not present in the tree, then 380 * we copy "data" into a new tree entry and return that node, setting *isNew 381 * to true. 382 * 383 * If the value represented by "data" is already present, then we call the 384 * combiner function to merge data into the existing node, and return the 385 * existing node, setting *isNew to false. 386 * 387 * "data" is unmodified in either case; it's typically just a local 388 * variable in the caller. 389 */ 390 RBTNode * 391 rbt_insert(RBTree *rbt, const RBTNode *data, bool *isNew) 392 { 393 RBTNode *current, 394 *parent, 395 *x; 396 int cmp; 397 398 /* find where node belongs */ 399 current = rbt->root; 400 parent = NULL; 401 cmp = 0; /* just to prevent compiler warning */ 402 403 while (current != RBTNIL) 404 { 405 cmp = rbt->comparator(data, current, rbt->arg); 406 if (cmp == 0) 407 { 408 /* 409 * Found node with given key. Apply combiner. 410 */ 411 rbt->combiner(current, data, rbt->arg); 412 *isNew = false; 413 return current; 414 } 415 parent = current; 416 current = (cmp < 0) ? current->left : current->right; 417 } 418 419 /* 420 * Value is not present, so create a new node containing data. 421 */ 422 *isNew = true; 423 424 x = rbt->allocfunc(rbt->arg); 425 426 x->color = RBTRED; 427 428 x->left = RBTNIL; 429 x->right = RBTNIL; 430 x->parent = parent; 431 rbt_copy_data(rbt, x, data); 432 433 /* insert node in tree */ 434 if (parent) 435 { 436 if (cmp < 0) 437 parent->left = x; 438 else 439 parent->right = x; 440 } 441 else 442 { 443 rbt->root = x; 444 } 445 446 rbt_insert_fixup(rbt, x); 447 448 return x; 449 } 450 451 /********************************************************************** 452 * Deletion * 453 **********************************************************************/ 454 455 /* 456 * Maintain Red-Black tree balance after deleting a black node. 457 */ 458 static void 459 rbt_delete_fixup(RBTree *rbt, RBTNode *x) 460 { 461 /* 462 * x is always a black node. Initially, it is the former child of the 463 * deleted node. Each iteration of this loop moves it higher up in the 464 * tree. 465 */ 466 while (x != rbt->root && x->color == RBTBLACK) 467 { 468 /* 469 * Left and right cases are symmetric. Any nodes that are children of 470 * x have a black-height one less than the remainder of the nodes in 471 * the tree. We rotate and recolor nodes to move the problem up the 472 * tree: at some stage we'll either fix the problem, or reach the root 473 * (where the black-height is allowed to decrease). 474 */ 475 if (x == x->parent->left) 476 { 477 RBTNode *w = x->parent->right; 478 479 if (w->color == RBTRED) 480 { 481 w->color = RBTBLACK; 482 x->parent->color = RBTRED; 483 484 rbt_rotate_left(rbt, x->parent); 485 w = x->parent->right; 486 } 487 488 if (w->left->color == RBTBLACK && w->right->color == RBTBLACK) 489 { 490 w->color = RBTRED; 491 492 x = x->parent; 493 } 494 else 495 { 496 if (w->right->color == RBTBLACK) 497 { 498 w->left->color = RBTBLACK; 499 w->color = RBTRED; 500 501 rbt_rotate_right(rbt, w); 502 w = x->parent->right; 503 } 504 w->color = x->parent->color; 505 x->parent->color = RBTBLACK; 506 w->right->color = RBTBLACK; 507 508 rbt_rotate_left(rbt, x->parent); 509 x = rbt->root; /* Arrange for loop to terminate. */ 510 } 511 } 512 else 513 { 514 RBTNode *w = x->parent->left; 515 516 if (w->color == RBTRED) 517 { 518 w->color = RBTBLACK; 519 x->parent->color = RBTRED; 520 521 rbt_rotate_right(rbt, x->parent); 522 w = x->parent->left; 523 } 524 525 if (w->right->color == RBTBLACK && w->left->color == RBTBLACK) 526 { 527 w->color = RBTRED; 528 529 x = x->parent; 530 } 531 else 532 { 533 if (w->left->color == RBTBLACK) 534 { 535 w->right->color = RBTBLACK; 536 w->color = RBTRED; 537 538 rbt_rotate_left(rbt, w); 539 w = x->parent->left; 540 } 541 w->color = x->parent->color; 542 x->parent->color = RBTBLACK; 543 w->left->color = RBTBLACK; 544 545 rbt_rotate_right(rbt, x->parent); 546 x = rbt->root; /* Arrange for loop to terminate. */ 547 } 548 } 549 } 550 x->color = RBTBLACK; 551 } 552 553 /* 554 * Delete node z from tree. 555 */ 556 static void 557 rbt_delete_node(RBTree *rbt, RBTNode *z) 558 { 559 RBTNode *x, 560 *y; 561 562 /* This is just paranoia: we should only get called on a valid node */ 563 if (!z || z == RBTNIL) 564 return; 565 566 /* 567 * y is the node that will actually be removed from the tree. This will 568 * be z if z has fewer than two children, or the tree successor of z 569 * otherwise. 570 */ 571 if (z->left == RBTNIL || z->right == RBTNIL) 572 { 573 /* y has a RBTNIL node as a child */ 574 y = z; 575 } 576 else 577 { 578 /* find tree successor */ 579 y = z->right; 580 while (y->left != RBTNIL) 581 y = y->left; 582 } 583 584 /* x is y's only child */ 585 if (y->left != RBTNIL) 586 x = y->left; 587 else 588 x = y->right; 589 590 /* Remove y from the tree. */ 591 x->parent = y->parent; 592 if (y->parent) 593 { 594 if (y == y->parent->left) 595 y->parent->left = x; 596 else 597 y->parent->right = x; 598 } 599 else 600 { 601 rbt->root = x; 602 } 603 604 /* 605 * If we removed the tree successor of z rather than z itself, then move 606 * the data for the removed node to the one we were supposed to remove. 607 */ 608 if (y != z) 609 rbt_copy_data(rbt, z, y); 610 611 /* 612 * Removing a black node might make some paths from root to leaf contain 613 * fewer black nodes than others, or it might make two red nodes adjacent. 614 */ 615 if (y->color == RBTBLACK) 616 rbt_delete_fixup(rbt, x); 617 618 /* Now we can recycle the y node */ 619 if (rbt->freefunc) 620 rbt->freefunc(y, rbt->arg); 621 } 622 623 /* 624 * rbt_delete: remove the given tree entry 625 * 626 * "node" must have previously been found via rbt_find or rbt_leftmost. 627 * It is caller's responsibility to free any subsidiary data attached 628 * to the node before calling rbt_delete. (Do *not* try to push that 629 * responsibility off to the freefunc, as some other physical node 630 * may be the one actually freed!) 631 */ 632 void 633 rbt_delete(RBTree *rbt, RBTNode *node) 634 { 635 rbt_delete_node(rbt, node); 636 } 637 638 /********************************************************************** 639 * Traverse * 640 **********************************************************************/ 641 642 static RBTNode * 643 rbt_left_right_iterator(RBTreeIterator *iter) 644 { 645 if (iter->last_visited == NULL) 646 { 647 iter->last_visited = iter->rbt->root; 648 while (iter->last_visited->left != RBTNIL) 649 iter->last_visited = iter->last_visited->left; 650 651 return iter->last_visited; 652 } 653 654 if (iter->last_visited->right != RBTNIL) 655 { 656 iter->last_visited = iter->last_visited->right; 657 while (iter->last_visited->left != RBTNIL) 658 iter->last_visited = iter->last_visited->left; 659 660 return iter->last_visited; 661 } 662 663 for (;;) 664 { 665 RBTNode *came_from = iter->last_visited; 666 667 iter->last_visited = iter->last_visited->parent; 668 if (iter->last_visited == NULL) 669 { 670 iter->is_over = true; 671 break; 672 } 673 674 if (iter->last_visited->left == came_from) 675 break; /* came from left sub-tree, return current 676 * node */ 677 678 /* else - came from right sub-tree, continue to move up */ 679 } 680 681 return iter->last_visited; 682 } 683 684 static RBTNode * 685 rbt_right_left_iterator(RBTreeIterator *iter) 686 { 687 if (iter->last_visited == NULL) 688 { 689 iter->last_visited = iter->rbt->root; 690 while (iter->last_visited->right != RBTNIL) 691 iter->last_visited = iter->last_visited->right; 692 693 return iter->last_visited; 694 } 695 696 if (iter->last_visited->left != RBTNIL) 697 { 698 iter->last_visited = iter->last_visited->left; 699 while (iter->last_visited->right != RBTNIL) 700 iter->last_visited = iter->last_visited->right; 701 702 return iter->last_visited; 703 } 704 705 for (;;) 706 { 707 RBTNode *came_from = iter->last_visited; 708 709 iter->last_visited = iter->last_visited->parent; 710 if (iter->last_visited == NULL) 711 { 712 iter->is_over = true; 713 break; 714 } 715 716 if (iter->last_visited->right == came_from) 717 break; /* came from right sub-tree, return current 718 * node */ 719 720 /* else - came from left sub-tree, continue to move up */ 721 } 722 723 return iter->last_visited; 724 } 725 726 /* 727 * rbt_begin_iterate: prepare to traverse the tree in any of several orders 728 * 729 * After calling rbt_begin_iterate, call rbt_iterate repeatedly until it 730 * returns NULL or the traversal stops being of interest. 731 * 732 * If the tree is changed during traversal, results of further calls to 733 * rbt_iterate are unspecified. Multiple concurrent iterators on the same 734 * tree are allowed. 735 * 736 * The iterator state is stored in the 'iter' struct. The caller should 737 * treat it as an opaque struct. 738 */ 739 void 740 rbt_begin_iterate(RBTree *rbt, RBTOrderControl ctrl, RBTreeIterator *iter) 741 { 742 /* Common initialization for all traversal orders */ 743 iter->rbt = rbt; 744 iter->last_visited = NULL; 745 iter->is_over = (rbt->root == RBTNIL); 746 747 switch (ctrl) 748 { 749 case LeftRightWalk: /* visit left, then self, then right */ 750 iter->iterate = rbt_left_right_iterator; 751 break; 752 case RightLeftWalk: /* visit right, then self, then left */ 753 iter->iterate = rbt_right_left_iterator; 754 break; 755 default: 756 elog(ERROR, "unrecognized rbtree iteration order: %d", ctrl); 757 } 758 } 759 760 /* 761 * rbt_iterate: return the next node in traversal order, or NULL if no more 762 */ 763 RBTNode * 764 rbt_iterate(RBTreeIterator *iter) 765 { 766 if (iter->is_over) 767 return NULL; 768 769 return iter->iterate(iter); 770 } 771