1/** @file 2 An OrderedCollectionLib instance that provides a red-black tree 3 implementation, and allocates and releases tree nodes with 4 MemoryAllocationLib. 5 6 This library instance is useful when a fast associative container is needed. 7 Worst case time complexity is O(log n) for Find(), Next(), Prev(), Min(), 8 Max(), Insert(), and Delete(), where "n" is the number of elements in the 9 tree. Complete ordered traversal takes O(n) time. 10 11 The implementation is also useful as a fast priority queue. 12 13 Copyright (C) 2014, Red Hat, Inc. 14 Copyright (c) 2014, Intel Corporation. All rights reserved.<BR> 15 16 This program and the accompanying materials are licensed and made available 17 under the terms and conditions of the BSD License that accompanies this 18 distribution. The full text of the license may be found at 19 http://opensource.org/licenses/bsd-license.php. 20 21 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS, WITHOUT 22 WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED. 23**/ 24 25#include <Library/OrderedCollectionLib.h> 26#include <Library/DebugLib.h> 27#include <Library/MemoryAllocationLib.h> 28 29typedef enum { 30 RedBlackTreeRed, 31 RedBlackTreeBlack 32} RED_BLACK_TREE_COLOR; 33 34// 35// Incomplete types and convenience typedefs are present in the library class 36// header. Beside completing the types, we introduce typedefs here that reflect 37// the implementation closely. 38// 39typedef ORDERED_COLLECTION RED_BLACK_TREE; 40typedef ORDERED_COLLECTION_ENTRY RED_BLACK_TREE_NODE; 41typedef ORDERED_COLLECTION_USER_COMPARE RED_BLACK_TREE_USER_COMPARE; 42typedef ORDERED_COLLECTION_KEY_COMPARE RED_BLACK_TREE_KEY_COMPARE; 43 44struct ORDERED_COLLECTION { 45 RED_BLACK_TREE_NODE *Root; 46 RED_BLACK_TREE_USER_COMPARE UserStructCompare; 47 RED_BLACK_TREE_KEY_COMPARE KeyCompare; 48}; 49 50struct ORDERED_COLLECTION_ENTRY { 51 VOID *UserStruct; 52 RED_BLACK_TREE_NODE *Parent; 53 RED_BLACK_TREE_NODE *Left; 54 RED_BLACK_TREE_NODE *Right; 55 RED_BLACK_TREE_COLOR Color; 56}; 57 58 59/** 60 Retrieve the user structure linked by the specified tree node. 61 62 Read-only operation. 63 64 @param[in] Node Pointer to the tree node whose associated user structure we 65 want to retrieve. The caller is responsible for passing a 66 non-NULL argument. 67 68 @return Pointer to user structure linked by Node. 69**/ 70VOID * 71EFIAPI 72OrderedCollectionUserStruct ( 73 IN CONST RED_BLACK_TREE_NODE *Node 74 ) 75{ 76 return Node->UserStruct; 77} 78 79/** 80 A slow function that asserts that the tree is a valid red-black tree, and 81 that it orders user structures correctly. 82 83 Read-only operation. 84 85 This function uses the stack for recursion and is not recommended for 86 "production use". 87 88 @param[in] Tree The tree to validate. 89**/ 90VOID 91RedBlackTreeValidate ( 92 IN CONST RED_BLACK_TREE *Tree 93 ); 94 95 96/** 97 Allocate and initialize the RED_BLACK_TREE structure. 98 99 Allocation occurs via MemoryAllocationLib's AllocatePool() function. 100 101 @param[in] UserStructCompare This caller-provided function will be used to 102 order two user structures linked into the 103 tree, during the insertion procedure. 104 105 @param[in] KeyCompare This caller-provided function will be used to 106 order the standalone search key against user 107 structures linked into the tree, during the 108 lookup procedure. 109 110 @retval NULL If allocation failed. 111 112 @return Pointer to the allocated, initialized RED_BLACK_TREE structure, 113 otherwise. 114**/ 115RED_BLACK_TREE * 116EFIAPI 117OrderedCollectionInit ( 118 IN RED_BLACK_TREE_USER_COMPARE UserStructCompare, 119 IN RED_BLACK_TREE_KEY_COMPARE KeyCompare 120 ) 121{ 122 RED_BLACK_TREE *Tree; 123 124 Tree = AllocatePool (sizeof *Tree); 125 if (Tree == NULL) { 126 return NULL; 127 } 128 129 Tree->Root = NULL; 130 Tree->UserStructCompare = UserStructCompare; 131 Tree->KeyCompare = KeyCompare; 132 133 if (FeaturePcdGet (PcdValidateOrderedCollection)) { 134 RedBlackTreeValidate (Tree); 135 } 136 return Tree; 137} 138 139 140/** 141 Check whether the tree is empty (has no nodes). 142 143 Read-only operation. 144 145 @param[in] Tree The tree to check for emptiness. 146 147 @retval TRUE The tree is empty. 148 149 @retval FALSE The tree is not empty. 150**/ 151BOOLEAN 152EFIAPI 153OrderedCollectionIsEmpty ( 154 IN CONST RED_BLACK_TREE *Tree 155 ) 156{ 157 return (BOOLEAN)(Tree->Root == NULL); 158} 159 160 161/** 162 Uninitialize and release an empty RED_BLACK_TREE structure. 163 164 Read-write operation. 165 166 Release occurs via MemoryAllocationLib's FreePool() function. 167 168 It is the caller's responsibility to delete all nodes from the tree before 169 calling this function. 170 171 @param[in] Tree The empty tree to uninitialize and release. 172**/ 173VOID 174EFIAPI 175OrderedCollectionUninit ( 176 IN RED_BLACK_TREE *Tree 177 ) 178{ 179 ASSERT (OrderedCollectionIsEmpty (Tree)); 180 FreePool (Tree); 181} 182 183 184/** 185 Look up the tree node that links the user structure that matches the 186 specified standalone key. 187 188 Read-only operation. 189 190 @param[in] Tree The tree to search for StandaloneKey. 191 192 @param[in] StandaloneKey The key to locate among the user structures linked 193 into Tree. StandaloneKey will be passed to 194 Tree->KeyCompare(). 195 196 @retval NULL StandaloneKey could not be found. 197 198 @return The tree node that links to the user structure matching 199 StandaloneKey, otherwise. 200**/ 201RED_BLACK_TREE_NODE * 202EFIAPI 203OrderedCollectionFind ( 204 IN CONST RED_BLACK_TREE *Tree, 205 IN CONST VOID *StandaloneKey 206 ) 207{ 208 RED_BLACK_TREE_NODE *Node; 209 210 Node = Tree->Root; 211 while (Node != NULL) { 212 INTN Result; 213 214 Result = Tree->KeyCompare (StandaloneKey, Node->UserStruct); 215 if (Result == 0) { 216 break; 217 } 218 Node = (Result < 0) ? Node->Left : Node->Right; 219 } 220 return Node; 221} 222 223 224/** 225 Find the tree node of the minimum user structure stored in the tree. 226 227 Read-only operation. 228 229 @param[in] Tree The tree to return the minimum node of. The user structure 230 linked by the minimum node compares less than all other user 231 structures in the tree. 232 233 @retval NULL If Tree is empty. 234 235 @return The tree node that links the minimum user structure, otherwise. 236**/ 237RED_BLACK_TREE_NODE * 238EFIAPI 239OrderedCollectionMin ( 240 IN CONST RED_BLACK_TREE *Tree 241 ) 242{ 243 RED_BLACK_TREE_NODE *Node; 244 245 Node = Tree->Root; 246 if (Node == NULL) { 247 return NULL; 248 } 249 while (Node->Left != NULL) { 250 Node = Node->Left; 251 } 252 return Node; 253} 254 255 256/** 257 Find the tree node of the maximum user structure stored in the tree. 258 259 Read-only operation. 260 261 @param[in] Tree The tree to return the maximum node of. The user structure 262 linked by the maximum node compares greater than all other 263 user structures in the tree. 264 265 @retval NULL If Tree is empty. 266 267 @return The tree node that links the maximum user structure, otherwise. 268**/ 269RED_BLACK_TREE_NODE * 270EFIAPI 271OrderedCollectionMax ( 272 IN CONST RED_BLACK_TREE *Tree 273 ) 274{ 275 RED_BLACK_TREE_NODE *Node; 276 277 Node = Tree->Root; 278 if (Node == NULL) { 279 return NULL; 280 } 281 while (Node->Right != NULL) { 282 Node = Node->Right; 283 } 284 return Node; 285} 286 287 288/** 289 Get the tree node of the least user structure that is greater than the one 290 linked by Node. 291 292 Read-only operation. 293 294 @param[in] Node The node to get the successor node of. 295 296 @retval NULL If Node is NULL, or Node is the maximum node of its containing 297 tree (ie. Node has no successor node). 298 299 @return The tree node linking the least user structure that is greater 300 than the one linked by Node, otherwise. 301**/ 302RED_BLACK_TREE_NODE * 303EFIAPI 304OrderedCollectionNext ( 305 IN CONST RED_BLACK_TREE_NODE *Node 306 ) 307{ 308 RED_BLACK_TREE_NODE *Walk; 309 CONST RED_BLACK_TREE_NODE *Child; 310 311 if (Node == NULL) { 312 return NULL; 313 } 314 315 // 316 // If Node has a right subtree, then the successor is the minimum node of 317 // that subtree. 318 // 319 Walk = Node->Right; 320 if (Walk != NULL) { 321 while (Walk->Left != NULL) { 322 Walk = Walk->Left; 323 } 324 return Walk; 325 } 326 327 // 328 // Otherwise we have to ascend as long as we're our parent's right child (ie. 329 // ascending to the left). 330 // 331 Child = Node; 332 Walk = Child->Parent; 333 while (Walk != NULL && Child == Walk->Right) { 334 Child = Walk; 335 Walk = Child->Parent; 336 } 337 return Walk; 338} 339 340 341/** 342 Get the tree node of the greatest user structure that is less than the one 343 linked by Node. 344 345 Read-only operation. 346 347 @param[in] Node The node to get the predecessor node of. 348 349 @retval NULL If Node is NULL, or Node is the minimum node of its containing 350 tree (ie. Node has no predecessor node). 351 352 @return The tree node linking the greatest user structure that is less 353 than the one linked by Node, otherwise. 354**/ 355RED_BLACK_TREE_NODE * 356EFIAPI 357OrderedCollectionPrev ( 358 IN CONST RED_BLACK_TREE_NODE *Node 359 ) 360{ 361 RED_BLACK_TREE_NODE *Walk; 362 CONST RED_BLACK_TREE_NODE *Child; 363 364 if (Node == NULL) { 365 return NULL; 366 } 367 368 // 369 // If Node has a left subtree, then the predecessor is the maximum node of 370 // that subtree. 371 // 372 Walk = Node->Left; 373 if (Walk != NULL) { 374 while (Walk->Right != NULL) { 375 Walk = Walk->Right; 376 } 377 return Walk; 378 } 379 380 // 381 // Otherwise we have to ascend as long as we're our parent's left child (ie. 382 // ascending to the right). 383 // 384 Child = Node; 385 Walk = Child->Parent; 386 while (Walk != NULL && Child == Walk->Left) { 387 Child = Walk; 388 Walk = Child->Parent; 389 } 390 return Walk; 391} 392 393 394/** 395 Rotate tree nodes around Pivot to the right. 396 397 Parent Parent 398 | | 399 Pivot LeftChild 400 / . . \_ 401 LeftChild Node1 ---> Node2 Pivot 402 . \ / . 403 Node2 LeftRightChild LeftRightChild Node1 404 405 The ordering Node2 < LeftChild < LeftRightChild < Pivot < Node1 is kept 406 intact. Parent (if any) is either at the left extreme or the right extreme of 407 this ordering, and that relation is also kept intact. 408 409 Edges marked with a dot (".") don't change during rotation. 410 411 Internal read-write operation. 412 413 @param[in,out] Pivot The tree node to rotate other nodes right around. It 414 is the caller's responsibility to ensure that 415 Pivot->Left is not NULL. 416 417 @param[out] NewRoot If Pivot has a parent node on input, then the 418 function updates Pivot's original parent on output 419 according to the rotation, and NewRoot is not 420 accessed. 421 422 If Pivot has no parent node on input (ie. Pivot is 423 the root of the tree), then the function stores the 424 new root node of the tree in NewRoot. 425**/ 426VOID 427RedBlackTreeRotateRight ( 428 IN OUT RED_BLACK_TREE_NODE *Pivot, 429 OUT RED_BLACK_TREE_NODE **NewRoot 430 ) 431{ 432 RED_BLACK_TREE_NODE *Parent; 433 RED_BLACK_TREE_NODE *LeftChild; 434 RED_BLACK_TREE_NODE *LeftRightChild; 435 436 Parent = Pivot->Parent; 437 LeftChild = Pivot->Left; 438 LeftRightChild = LeftChild->Right; 439 440 Pivot->Left = LeftRightChild; 441 if (LeftRightChild != NULL) { 442 LeftRightChild->Parent = Pivot; 443 } 444 LeftChild->Parent = Parent; 445 if (Parent == NULL) { 446 *NewRoot = LeftChild; 447 } else { 448 if (Pivot == Parent->Left) { 449 Parent->Left = LeftChild; 450 } else { 451 Parent->Right = LeftChild; 452 } 453 } 454 LeftChild->Right = Pivot; 455 Pivot->Parent = LeftChild; 456} 457 458 459/** 460 Rotate tree nodes around Pivot to the left. 461 462 Parent Parent 463 | | 464 Pivot RightChild 465 . \ / . 466 Node1 RightChild ---> Pivot Node2 467 /. . \_ 468 RightLeftChild Node2 Node1 RightLeftChild 469 470 The ordering Node1 < Pivot < RightLeftChild < RightChild < Node2 is kept 471 intact. Parent (if any) is either at the left extreme or the right extreme of 472 this ordering, and that relation is also kept intact. 473 474 Edges marked with a dot (".") don't change during rotation. 475 476 Internal read-write operation. 477 478 @param[in,out] Pivot The tree node to rotate other nodes left around. It 479 is the caller's responsibility to ensure that 480 Pivot->Right is not NULL. 481 482 @param[out] NewRoot If Pivot has a parent node on input, then the 483 function updates Pivot's original parent on output 484 according to the rotation, and NewRoot is not 485 accessed. 486 487 If Pivot has no parent node on input (ie. Pivot is 488 the root of the tree), then the function stores the 489 new root node of the tree in NewRoot. 490**/ 491VOID 492RedBlackTreeRotateLeft ( 493 IN OUT RED_BLACK_TREE_NODE *Pivot, 494 OUT RED_BLACK_TREE_NODE **NewRoot 495 ) 496{ 497 RED_BLACK_TREE_NODE *Parent; 498 RED_BLACK_TREE_NODE *RightChild; 499 RED_BLACK_TREE_NODE *RightLeftChild; 500 501 Parent = Pivot->Parent; 502 RightChild = Pivot->Right; 503 RightLeftChild = RightChild->Left; 504 505 Pivot->Right = RightLeftChild; 506 if (RightLeftChild != NULL) { 507 RightLeftChild->Parent = Pivot; 508 } 509 RightChild->Parent = Parent; 510 if (Parent == NULL) { 511 *NewRoot = RightChild; 512 } else { 513 if (Pivot == Parent->Left) { 514 Parent->Left = RightChild; 515 } else { 516 Parent->Right = RightChild; 517 } 518 } 519 RightChild->Left = Pivot; 520 Pivot->Parent = RightChild; 521} 522 523 524/** 525 Insert (link) a user structure into the tree. 526 527 Read-write operation. 528 529 This function allocates the new tree node with MemoryAllocationLib's 530 AllocatePool() function. 531 532 @param[in,out] Tree The tree to insert UserStruct into. 533 534 @param[out] Node The meaning of this optional, output-only 535 parameter depends on the return value of the 536 function. 537 538 When insertion is successful (RETURN_SUCCESS), 539 Node is set on output to the new tree node that 540 now links UserStruct. 541 542 When insertion fails due to lack of memory 543 (RETURN_OUT_OF_RESOURCES), Node is not changed. 544 545 When insertion fails due to key collision (ie. 546 another user structure is already in the tree that 547 compares equal to UserStruct), with return value 548 RETURN_ALREADY_STARTED, then Node is set on output 549 to the node that links the colliding user 550 structure. This enables "find-or-insert" in one 551 function call, or helps with later removal of the 552 colliding element. 553 554 @param[in] UserStruct The user structure to link into the tree. 555 UserStruct is ordered against in-tree user 556 structures with the Tree->UserStructCompare() 557 function. 558 559 @retval RETURN_SUCCESS Insertion successful. A new tree node has 560 been allocated, linking UserStruct. The new 561 tree node is reported back in Node (if the 562 caller requested it). 563 564 Existing RED_BLACK_TREE_NODE pointers into 565 Tree remain valid. For example, on-going 566 iterations in the caller can continue with 567 OrderedCollectionNext() / 568 OrderedCollectionPrev(), and they will 569 return the new node at some point if user 570 structure order dictates it. 571 572 @retval RETURN_OUT_OF_RESOURCES AllocatePool() failed to allocate memory for 573 the new tree node. The tree has not been 574 changed. Existing RED_BLACK_TREE_NODE 575 pointers into Tree remain valid. 576 577 @retval RETURN_ALREADY_STARTED A user structure has been found in the tree 578 that compares equal to UserStruct. The node 579 linking the colliding user structure is 580 reported back in Node (if the caller 581 requested it). The tree has not been 582 changed. Existing RED_BLACK_TREE_NODE 583 pointers into Tree remain valid. 584**/ 585RETURN_STATUS 586EFIAPI 587OrderedCollectionInsert ( 588 IN OUT RED_BLACK_TREE *Tree, 589 OUT RED_BLACK_TREE_NODE **Node OPTIONAL, 590 IN VOID *UserStruct 591 ) 592{ 593 RED_BLACK_TREE_NODE *Tmp; 594 RED_BLACK_TREE_NODE *Parent; 595 INTN Result; 596 RETURN_STATUS Status; 597 RED_BLACK_TREE_NODE *NewRoot; 598 599 Tmp = Tree->Root; 600 Parent = NULL; 601 Result = 0; 602 603 // 604 // First look for a collision, saving the last examined node for the case 605 // when there's no collision. 606 // 607 while (Tmp != NULL) { 608 Result = Tree->UserStructCompare (UserStruct, Tmp->UserStruct); 609 if (Result == 0) { 610 break; 611 } 612 Parent = Tmp; 613 Tmp = (Result < 0) ? Tmp->Left : Tmp->Right; 614 } 615 616 if (Tmp != NULL) { 617 if (Node != NULL) { 618 *Node = Tmp; 619 } 620 Status = RETURN_ALREADY_STARTED; 621 goto Done; 622 } 623 624 // 625 // no collision, allocate a new node 626 // 627 Tmp = AllocatePool (sizeof *Tmp); 628 if (Tmp == NULL) { 629 Status = RETURN_OUT_OF_RESOURCES; 630 goto Done; 631 } 632 if (Node != NULL) { 633 *Node = Tmp; 634 } 635 636 // 637 // reference the user structure from the node 638 // 639 Tmp->UserStruct = UserStruct; 640 641 // 642 // Link the node as a child to the correct side of the parent. 643 // If there's no parent, the new node is the root node in the tree. 644 // 645 Tmp->Parent = Parent; 646 Tmp->Left = NULL; 647 Tmp->Right = NULL; 648 if (Parent == NULL) { 649 Tree->Root = Tmp; 650 Tmp->Color = RedBlackTreeBlack; 651 Status = RETURN_SUCCESS; 652 goto Done; 653 } 654 if (Result < 0) { 655 Parent->Left = Tmp; 656 } else { 657 Parent->Right = Tmp; 658 } 659 Tmp->Color = RedBlackTreeRed; 660 661 // 662 // Red-black tree properties: 663 // 664 // #1 Each node is either red or black (RED_BLACK_TREE_NODE.Color). 665 // 666 // #2 Each leaf (ie. a pseudo-node pointed-to by a NULL valued 667 // RED_BLACK_TREE_NODE.Left or RED_BLACK_TREE_NODE.Right field) is black. 668 // 669 // #3 Each red node has two black children. 670 // 671 // #4 For any node N, and for any leaves L1 and L2 reachable from N, the 672 // paths N..L1 and N..L2 contain the same number of black nodes. 673 // 674 // #5 The root node is black. 675 // 676 // By replacing a leaf with a red node above, only property #3 may have been 677 // broken. (Note that this is the only edge across which property #3 might 678 // not hold in the entire tree.) Restore property #3. 679 // 680 681 NewRoot = Tree->Root; 682 while (Tmp != NewRoot && Parent->Color == RedBlackTreeRed) { 683 RED_BLACK_TREE_NODE *GrandParent; 684 RED_BLACK_TREE_NODE *Uncle; 685 686 // 687 // Tmp is not the root node. Tmp is red. Tmp's parent is red. (Breaking 688 // property #3.) 689 // 690 // Due to property #5, Tmp's parent cannot be the root node, hence Tmp's 691 // grandparent exists. 692 // 693 // Tmp's grandparent is black, because property #3 is only broken between 694 // Tmp and Tmp's parent. 695 // 696 GrandParent = Parent->Parent; 697 698 if (Parent == GrandParent->Left) { 699 Uncle = GrandParent->Right; 700 if (Uncle != NULL && Uncle->Color == RedBlackTreeRed) { 701 // 702 // GrandParent (black) 703 // / \_ 704 // Parent (red) Uncle (red) 705 // | 706 // Tmp (red) 707 // 708 709 Parent->Color = RedBlackTreeBlack; 710 Uncle->Color = RedBlackTreeBlack; 711 GrandParent->Color = RedBlackTreeRed; 712 713 // 714 // GrandParent (red) 715 // / \_ 716 // Parent (black) Uncle (black) 717 // | 718 // Tmp (red) 719 // 720 // We restored property #3 between Tmp and Tmp's parent, without 721 // breaking property #4. However, we may have broken property #3 722 // between Tmp's grandparent and Tmp's great-grandparent (if any), so 723 // repeat the loop for Tmp's grandparent. 724 // 725 // If Tmp's grandparent has no parent, then the loop will terminate, 726 // and we will have broken property #5, by coloring the root red. We'll 727 // restore property #5 after the loop, without breaking any others. 728 // 729 Tmp = GrandParent; 730 Parent = Tmp->Parent; 731 } else { 732 // 733 // Tmp's uncle is black (satisfied by the case too when Tmp's uncle is 734 // NULL, see property #2). 735 // 736 737 if (Tmp == Parent->Right) { 738 // 739 // GrandParent (black): D 740 // / \_ 741 // Parent (red): A Uncle (black): E 742 // \_ 743 // Tmp (red): B 744 // \_ 745 // black: C 746 // 747 // Rotate left, pivoting on node A. This keeps the breakage of 748 // property #3 in the same spot, and keeps other properties intact 749 // (because both Tmp and its parent are red). 750 // 751 Tmp = Parent; 752 RedBlackTreeRotateLeft (Tmp, &NewRoot); 753 Parent = Tmp->Parent; 754 755 // 756 // With the rotation we reached the same configuration as if Tmp had 757 // been a left child to begin with. 758 // 759 // GrandParent (black): D 760 // / \_ 761 // Parent (red): B Uncle (black): E 762 // / \_ 763 // Tmp (red): A black: C 764 // 765 ASSERT (GrandParent == Parent->Parent); 766 } 767 768 Parent->Color = RedBlackTreeBlack; 769 GrandParent->Color = RedBlackTreeRed; 770 771 // 772 // Property #3 is now restored, but we've broken property #4. Namely, 773 // paths going through node E now see a decrease in black count, while 774 // paths going through node B don't. 775 // 776 // GrandParent (red): D 777 // / \_ 778 // Parent (black): B Uncle (black): E 779 // / \_ 780 // Tmp (red): A black: C 781 // 782 783 RedBlackTreeRotateRight (GrandParent, &NewRoot); 784 785 // 786 // Property #4 has been restored for node E, and preserved for others. 787 // 788 // Parent (black): B 789 // / \_ 790 // Tmp (red): A [GrandParent] (red): D 791 // / \_ 792 // black: C [Uncle] (black): E 793 // 794 // This configuration terminates the loop because Tmp's parent is now 795 // black. 796 // 797 } 798 } else { 799 // 800 // Symmetrical to the other branch. 801 // 802 Uncle = GrandParent->Left; 803 if (Uncle != NULL && Uncle->Color == RedBlackTreeRed) { 804 Parent->Color = RedBlackTreeBlack; 805 Uncle->Color = RedBlackTreeBlack; 806 GrandParent->Color = RedBlackTreeRed; 807 Tmp = GrandParent; 808 Parent = Tmp->Parent; 809 } else { 810 if (Tmp == Parent->Left) { 811 Tmp = Parent; 812 RedBlackTreeRotateRight (Tmp, &NewRoot); 813 Parent = Tmp->Parent; 814 ASSERT (GrandParent == Parent->Parent); 815 } 816 Parent->Color = RedBlackTreeBlack; 817 GrandParent->Color = RedBlackTreeRed; 818 RedBlackTreeRotateLeft (GrandParent, &NewRoot); 819 } 820 } 821 } 822 823 NewRoot->Color = RedBlackTreeBlack; 824 Tree->Root = NewRoot; 825 Status = RETURN_SUCCESS; 826 827Done: 828 if (FeaturePcdGet (PcdValidateOrderedCollection)) { 829 RedBlackTreeValidate (Tree); 830 } 831 return Status; 832} 833 834 835/** 836 Check if a node is black, allowing for leaf nodes (see property #2). 837 838 This is a convenience shorthand. 839 840 param[in] Node The node to check. Node may be NULL, corresponding to a leaf. 841 842 @return If Node is NULL or colored black. 843**/ 844BOOLEAN 845NodeIsNullOrBlack ( 846 IN CONST RED_BLACK_TREE_NODE *Node 847 ) 848{ 849 return (BOOLEAN)(Node == NULL || Node->Color == RedBlackTreeBlack); 850} 851 852 853/** 854 Delete a node from the tree, unlinking the associated user structure. 855 856 Read-write operation. 857 858 @param[in,out] Tree The tree to delete Node from. 859 860 @param[in] Node The tree node to delete from Tree. The caller is 861 responsible for ensuring that Node belongs to 862 Tree, and that Node is non-NULL and valid. Node is 863 typically an earlier return value, or output 864 parameter, of: 865 866 - OrderedCollectionFind(), for deleting a node by 867 user structure key, 868 869 - OrderedCollectionMin() / OrderedCollectionMax(), 870 for deleting the minimum / maximum node, 871 872 - OrderedCollectionNext() / 873 OrderedCollectionPrev(), for deleting a node 874 found during an iteration, 875 876 - OrderedCollectionInsert() with return value 877 RETURN_ALREADY_STARTED, for deleting a node 878 whose linked user structure caused collision 879 during insertion. 880 881 Given a non-empty Tree, Tree->Root is also a valid 882 Node argument (typically used for simplicity in 883 loops that empty the tree completely). 884 885 Node is released with MemoryAllocationLib's 886 FreePool() function. 887 888 Existing RED_BLACK_TREE_NODE pointers (ie. 889 iterators) *different* from Node remain valid. For 890 example: 891 892 - OrderedCollectionNext() / 893 OrderedCollectionPrev() iterations in the caller 894 can be continued from Node, if 895 OrderedCollectionNext() or 896 OrderedCollectionPrev() is called on Node 897 *before* OrderedCollectionDelete() is. That is, 898 fetch the successor / predecessor node first, 899 then delete Node. 900 901 - On-going iterations in the caller that would 902 have otherwise returned Node at some point, as 903 dictated by user structure order, will correctly 904 reflect the absence of Node after 905 OrderedCollectionDelete() is called 906 mid-iteration. 907 908 @param[out] UserStruct If the caller provides this optional output-only 909 parameter, then on output it is set to the user 910 structure originally linked by Node (which is now 911 freed). 912 913 This is a convenience that may save the caller a 914 OrderedCollectionUserStruct() invocation before 915 calling OrderedCollectionDelete(), in order to 916 retrieve the user structure being unlinked. 917**/ 918VOID 919EFIAPI 920OrderedCollectionDelete ( 921 IN OUT RED_BLACK_TREE *Tree, 922 IN RED_BLACK_TREE_NODE *Node, 923 OUT VOID **UserStruct OPTIONAL 924 ) 925{ 926 RED_BLACK_TREE_NODE *NewRoot; 927 RED_BLACK_TREE_NODE *OrigLeftChild; 928 RED_BLACK_TREE_NODE *OrigRightChild; 929 RED_BLACK_TREE_NODE *OrigParent; 930 RED_BLACK_TREE_NODE *Child; 931 RED_BLACK_TREE_NODE *Parent; 932 RED_BLACK_TREE_COLOR ColorOfUnlinked; 933 934 NewRoot = Tree->Root; 935 OrigLeftChild = Node->Left, 936 OrigRightChild = Node->Right, 937 OrigParent = Node->Parent; 938 939 if (UserStruct != NULL) { 940 *UserStruct = Node->UserStruct; 941 } 942 943 // 944 // After this block, no matter which branch we take: 945 // - Child will point to the unique (or NULL) original child of the node that 946 // we will have unlinked, 947 // - Parent will point to the *position* of the original parent of the node 948 // that we will have unlinked. 949 // 950 if (OrigLeftChild == NULL || OrigRightChild == NULL) { 951 // 952 // Node has at most one child. We can connect that child (if any) with 953 // Node's parent (if any), unlinking Node. This will preserve ordering 954 // because the subtree rooted in Node's child (if any) remains on the same 955 // side of Node's parent (if any) that Node was before. 956 // 957 Parent = OrigParent; 958 Child = (OrigLeftChild != NULL) ? OrigLeftChild : OrigRightChild; 959 ColorOfUnlinked = Node->Color; 960 961 if (Child != NULL) { 962 Child->Parent = Parent; 963 } 964 if (OrigParent == NULL) { 965 NewRoot = Child; 966 } else { 967 if (Node == OrigParent->Left) { 968 OrigParent->Left = Child; 969 } else { 970 OrigParent->Right = Child; 971 } 972 } 973 } else { 974 // 975 // Node has two children. We unlink Node's successor, and then link it into 976 // Node's place, keeping Node's original color. This preserves ordering 977 // because: 978 // - Node's left subtree is less than Node, hence less than Node's 979 // successor. 980 // - Node's right subtree is greater than Node. Node's successor is the 981 // minimum of that subtree, hence Node's successor is less than Node's 982 // right subtree with its minimum removed. 983 // - Node's successor is in Node's subtree, hence it falls on the same side 984 // of Node's parent as Node itself. The relinking doesn't change this 985 // relation. 986 // 987 RED_BLACK_TREE_NODE *ToRelink; 988 989 ToRelink = OrigRightChild; 990 if (ToRelink->Left == NULL) { 991 // 992 // OrigRightChild itself is Node's successor, it has no left child: 993 // 994 // OrigParent 995 // | 996 // Node: B 997 // / \_ 998 // OrigLeftChild: A OrigRightChild: E <--- Parent, ToRelink 999 // \_ 1000 // F <--- Child 1001 // 1002 Parent = OrigRightChild; 1003 Child = OrigRightChild->Right; 1004 } else { 1005 do { 1006 ToRelink = ToRelink->Left; 1007 } while (ToRelink->Left != NULL); 1008 1009 // 1010 // Node's successor is the minimum of OrigRightChild's proper subtree: 1011 // 1012 // OrigParent 1013 // | 1014 // Node: B 1015 // / \_ 1016 // OrigLeftChild: A OrigRightChild: E <--- Parent 1017 // / 1018 // C <--- ToRelink 1019 // \_ 1020 // D <--- Child 1021 Parent = ToRelink->Parent; 1022 Child = ToRelink->Right; 1023 1024 // 1025 // Unlink Node's successor (ie. ToRelink): 1026 // 1027 // OrigParent 1028 // | 1029 // Node: B 1030 // / \_ 1031 // OrigLeftChild: A OrigRightChild: E <--- Parent 1032 // / 1033 // D <--- Child 1034 // 1035 // C <--- ToRelink 1036 // 1037 Parent->Left = Child; 1038 if (Child != NULL) { 1039 Child->Parent = Parent; 1040 } 1041 1042 // 1043 // We start to link Node's unlinked successor into Node's place: 1044 // 1045 // OrigParent 1046 // | 1047 // Node: B C <--- ToRelink 1048 // / \_ 1049 // OrigLeftChild: A OrigRightChild: E <--- Parent 1050 // / 1051 // D <--- Child 1052 // 1053 // 1054 // 1055 ToRelink->Right = OrigRightChild; 1056 OrigRightChild->Parent = ToRelink; 1057 } 1058 1059 // 1060 // The rest handles both cases, attaching ToRelink (Node's original 1061 // successor) to OrigLeftChild and OrigParent. 1062 // 1063 // Parent, 1064 // OrigParent ToRelink OrigParent 1065 // | | | 1066 // Node: B | Node: B Parent 1067 // v | 1068 // OrigRightChild: E C <--- ToRelink | 1069 // / \ / \ v 1070 // OrigLeftChild: A F OrigLeftChild: A OrigRightChild: E 1071 // ^ / 1072 // | D <--- Child 1073 // Child 1074 // 1075 ToRelink->Left = OrigLeftChild; 1076 OrigLeftChild->Parent = ToRelink; 1077 1078 // 1079 // Node's color must be preserved in Node's original place. 1080 // 1081 ColorOfUnlinked = ToRelink->Color; 1082 ToRelink->Color = Node->Color; 1083 1084 // 1085 // Finish linking Node's unlinked successor into Node's place. 1086 // 1087 // Parent, 1088 // Node: B ToRelink Node: B 1089 // | 1090 // OrigParent | OrigParent Parent 1091 // | v | | 1092 // OrigRightChild: E C <--- ToRelink | 1093 // / \ / \ v 1094 // OrigLeftChild: A F OrigLeftChild: A OrigRightChild: E 1095 // ^ / 1096 // | D <--- Child 1097 // Child 1098 // 1099 ToRelink->Parent = OrigParent; 1100 if (OrigParent == NULL) { 1101 NewRoot = ToRelink; 1102 } else { 1103 if (Node == OrigParent->Left) { 1104 OrigParent->Left = ToRelink; 1105 } else { 1106 OrigParent->Right = ToRelink; 1107 } 1108 } 1109 } 1110 1111 FreePool (Node); 1112 1113 // 1114 // If the node that we unlinked from its original spot (ie. Node itself, or 1115 // Node's successor), was red, then we broke neither property #3 nor property 1116 // #4: we didn't create any red-red edge between Child and Parent, and we 1117 // didn't change the black count on any path. 1118 // 1119 if (ColorOfUnlinked == RedBlackTreeBlack) { 1120 // 1121 // However, if the unlinked node was black, then we have to transfer its 1122 // "black-increment" to its unique child (pointed-to by Child), lest we 1123 // break property #4 for its ancestors. 1124 // 1125 // If Child is red, we can simply color it black. If Child is black 1126 // already, we can't technically transfer a black-increment to it, due to 1127 // property #1. 1128 // 1129 // In the following loop we ascend searching for a red node to color black, 1130 // or until we reach the root (in which case we can drop the 1131 // black-increment). Inside the loop body, Child has a black value of 2, 1132 // transitorily breaking property #1 locally, but maintaining property #4 1133 // globally. 1134 // 1135 // Rotations in the loop preserve property #4. 1136 // 1137 while (Child != NewRoot && NodeIsNullOrBlack (Child)) { 1138 RED_BLACK_TREE_NODE *Sibling; 1139 RED_BLACK_TREE_NODE *LeftNephew; 1140 RED_BLACK_TREE_NODE *RightNephew; 1141 1142 if (Child == Parent->Left) { 1143 Sibling = Parent->Right; 1144 // 1145 // Sibling can never be NULL (ie. a leaf). 1146 // 1147 // If Sibling was NULL, then the black count on the path from Parent to 1148 // Sibling would equal Parent's black value, plus 1 (due to property 1149 // #2). Whereas the black count on the path from Parent to any leaf via 1150 // Child would be at least Parent's black value, plus 2 (due to Child's 1151 // black value of 2). This would clash with property #4. 1152 // 1153 // (Sibling can be black of course, but it has to be an internal node. 1154 // Internality allows Sibling to have children, bumping the black 1155 // counts of paths that go through it.) 1156 // 1157 ASSERT (Sibling != NULL); 1158 if (Sibling->Color == RedBlackTreeRed) { 1159 // 1160 // Sibling's red color implies its children (if any), node C and node 1161 // E, are black (property #3). It also implies that Parent is black. 1162 // 1163 // grandparent grandparent 1164 // | | 1165 // Parent,b:B b:D 1166 // / \ / \_ 1167 // Child,2b:A Sibling,r:D ---> Parent,r:B b:E 1168 // /\ /\_ 1169 // b:C b:E Child,2b:A Sibling,b:C 1170 // 1171 Sibling->Color = RedBlackTreeBlack; 1172 Parent->Color = RedBlackTreeRed; 1173 RedBlackTreeRotateLeft (Parent, &NewRoot); 1174 Sibling = Parent->Right; 1175 // 1176 // Same reasoning as above. 1177 // 1178 ASSERT (Sibling != NULL); 1179 } 1180 1181 // 1182 // Sibling is black, and not NULL. (Ie. Sibling is a black internal 1183 // node.) 1184 // 1185 ASSERT (Sibling->Color == RedBlackTreeBlack); 1186 LeftNephew = Sibling->Left; 1187 RightNephew = Sibling->Right; 1188 if (NodeIsNullOrBlack (LeftNephew) && 1189 NodeIsNullOrBlack (RightNephew)) { 1190 // 1191 // In this case we can "steal" one black value from Child and Sibling 1192 // each, and pass it to Parent. "Stealing" means that Sibling (black 1193 // value 1) becomes red, Child (black value 2) becomes singly-black, 1194 // and Parent will have to be examined if it can eat the 1195 // black-increment. 1196 // 1197 // Sibling is allowed to become red because both of its children are 1198 // black (property #3). 1199 // 1200 // grandparent Parent 1201 // | | 1202 // Parent,x:B Child,x:B 1203 // / \ / \_ 1204 // Child,2b:A Sibling,b:D ---> b:A r:D 1205 // /\ /\_ 1206 // LeftNephew,b:C RightNephew,b:E b:C b:E 1207 // 1208 Sibling->Color = RedBlackTreeRed; 1209 Child = Parent; 1210 Parent = Parent->Parent; 1211 // 1212 // Continue ascending. 1213 // 1214 } else { 1215 // 1216 // At least one nephew is red. 1217 // 1218 if (NodeIsNullOrBlack (RightNephew)) { 1219 // 1220 // Since the right nephew is black, the left nephew is red. Due to 1221 // property #3, LeftNephew has two black children, hence node E is 1222 // black. 1223 // 1224 // Together with the rotation, this enables us to color node F red 1225 // (because property #3 will be satisfied). We flip node D to black 1226 // to maintain property #4. 1227 // 1228 // grandparent grandparent 1229 // | | 1230 // Parent,x:B Parent,x:B 1231 // /\ /\_ 1232 // Child,2b:A Sibling,b:F ---> Child,2b:A Sibling,b:D 1233 // /\ / \_ 1234 // LeftNephew,r:D RightNephew,b:G b:C RightNephew,r:F 1235 // /\ /\_ 1236 // b:C b:E b:E b:G 1237 // 1238 LeftNephew->Color = RedBlackTreeBlack; 1239 Sibling->Color = RedBlackTreeRed; 1240 RedBlackTreeRotateRight (Sibling, &NewRoot); 1241 Sibling = Parent->Right; 1242 RightNephew = Sibling->Right; 1243 // 1244 // These operations ensure that... 1245 // 1246 } 1247 // 1248 // ... RightNephew is definitely red here, plus Sibling is (still) 1249 // black and non-NULL. 1250 // 1251 ASSERT (RightNephew != NULL); 1252 ASSERT (RightNephew->Color == RedBlackTreeRed); 1253 ASSERT (Sibling != NULL); 1254 ASSERT (Sibling->Color == RedBlackTreeBlack); 1255 // 1256 // In this case we can flush the extra black-increment immediately, 1257 // restoring property #1 for Child (node A): we color RightNephew 1258 // (node E) from red to black. 1259 // 1260 // In order to maintain property #4, we exchange colors between 1261 // Parent and Sibling (nodes B and D), and rotate left around Parent 1262 // (node B). The transformation doesn't change the black count 1263 // increase incurred by each partial path, eg. 1264 // - ascending from node A: 2 + x == 1 + 1 + x 1265 // - ascending from node C: y + 1 + x == y + 1 + x 1266 // - ascending from node E: 0 + 1 + x == 1 + x 1267 // 1268 // The color exchange is valid, because even if x stands for red, 1269 // both children of node D are black after the transformation 1270 // (preserving property #3). 1271 // 1272 // grandparent grandparent 1273 // | | 1274 // Parent,x:B x:D 1275 // / \ / \_ 1276 // Child,2b:A Sibling,b:D ---> b:B b:E 1277 // / \ / \_ 1278 // y:C RightNephew,r:E b:A y:C 1279 // 1280 // 1281 Sibling->Color = Parent->Color; 1282 Parent->Color = RedBlackTreeBlack; 1283 RightNephew->Color = RedBlackTreeBlack; 1284 RedBlackTreeRotateLeft (Parent, &NewRoot); 1285 Child = NewRoot; 1286 // 1287 // This terminates the loop. 1288 // 1289 } 1290 } else { 1291 // 1292 // Mirrors the other branch. 1293 // 1294 Sibling = Parent->Left; 1295 ASSERT (Sibling != NULL); 1296 if (Sibling->Color == RedBlackTreeRed) { 1297 Sibling->Color = RedBlackTreeBlack; 1298 Parent->Color = RedBlackTreeRed; 1299 RedBlackTreeRotateRight (Parent, &NewRoot); 1300 Sibling = Parent->Left; 1301 ASSERT (Sibling != NULL); 1302 } 1303 1304 ASSERT (Sibling->Color == RedBlackTreeBlack); 1305 RightNephew = Sibling->Right; 1306 LeftNephew = Sibling->Left; 1307 if (NodeIsNullOrBlack (RightNephew) && 1308 NodeIsNullOrBlack (LeftNephew)) { 1309 Sibling->Color = RedBlackTreeRed; 1310 Child = Parent; 1311 Parent = Parent->Parent; 1312 } else { 1313 if (NodeIsNullOrBlack (LeftNephew)) { 1314 RightNephew->Color = RedBlackTreeBlack; 1315 Sibling->Color = RedBlackTreeRed; 1316 RedBlackTreeRotateLeft (Sibling, &NewRoot); 1317 Sibling = Parent->Left; 1318 LeftNephew = Sibling->Left; 1319 } 1320 ASSERT (LeftNephew != NULL); 1321 ASSERT (LeftNephew->Color == RedBlackTreeRed); 1322 ASSERT (Sibling != NULL); 1323 ASSERT (Sibling->Color == RedBlackTreeBlack); 1324 Sibling->Color = Parent->Color; 1325 Parent->Color = RedBlackTreeBlack; 1326 LeftNephew->Color = RedBlackTreeBlack; 1327 RedBlackTreeRotateRight (Parent, &NewRoot); 1328 Child = NewRoot; 1329 } 1330 } 1331 } 1332 1333 if (Child != NULL) { 1334 Child->Color = RedBlackTreeBlack; 1335 } 1336 } 1337 1338 Tree->Root = NewRoot; 1339 1340 if (FeaturePcdGet (PcdValidateOrderedCollection)) { 1341 RedBlackTreeValidate (Tree); 1342 } 1343} 1344 1345 1346/** 1347 Recursively check the red-black tree properties #1 to #4 on a node. 1348 1349 @param[in] Node The root of the subtree to validate. 1350 1351 @retval The black-height of Node's parent. 1352**/ 1353UINT32 1354RedBlackTreeRecursiveCheck ( 1355 IN CONST RED_BLACK_TREE_NODE *Node 1356 ) 1357{ 1358 UINT32 LeftHeight; 1359 UINT32 RightHeight; 1360 1361 // 1362 // property #2 1363 // 1364 if (Node == NULL) { 1365 return 1; 1366 } 1367 1368 // 1369 // property #1 1370 // 1371 ASSERT (Node->Color == RedBlackTreeRed || Node->Color == RedBlackTreeBlack); 1372 1373 // 1374 // property #3 1375 // 1376 if (Node->Color == RedBlackTreeRed) { 1377 ASSERT (NodeIsNullOrBlack (Node->Left)); 1378 ASSERT (NodeIsNullOrBlack (Node->Right)); 1379 } 1380 1381 // 1382 // property #4 1383 // 1384 LeftHeight = RedBlackTreeRecursiveCheck (Node->Left); 1385 RightHeight = RedBlackTreeRecursiveCheck (Node->Right); 1386 ASSERT (LeftHeight == RightHeight); 1387 1388 return (Node->Color == RedBlackTreeBlack) + LeftHeight; 1389} 1390 1391 1392/** 1393 A slow function that asserts that the tree is a valid red-black tree, and 1394 that it orders user structures correctly. 1395 1396 Read-only operation. 1397 1398 This function uses the stack for recursion and is not recommended for 1399 "production use". 1400 1401 @param[in] Tree The tree to validate. 1402**/ 1403VOID 1404RedBlackTreeValidate ( 1405 IN CONST RED_BLACK_TREE *Tree 1406 ) 1407{ 1408 UINT32 BlackHeight; 1409 UINT32 ForwardCount; 1410 UINT32 BackwardCount; 1411 CONST RED_BLACK_TREE_NODE *Last; 1412 CONST RED_BLACK_TREE_NODE *Node; 1413 1414 DEBUG ((DEBUG_VERBOSE, "%a: Tree=%p\n", __FUNCTION__, Tree)); 1415 1416 // 1417 // property #5 1418 // 1419 ASSERT (NodeIsNullOrBlack (Tree->Root)); 1420 1421 // 1422 // check the other properties 1423 // 1424 BlackHeight = RedBlackTreeRecursiveCheck (Tree->Root) - 1; 1425 1426 // 1427 // forward ordering 1428 // 1429 Last = OrderedCollectionMin (Tree); 1430 ForwardCount = (Last != NULL); 1431 for (Node = OrderedCollectionNext (Last); Node != NULL; 1432 Node = OrderedCollectionNext (Last)) { 1433 ASSERT (Tree->UserStructCompare (Last->UserStruct, Node->UserStruct) < 0); 1434 Last = Node; 1435 ++ForwardCount; 1436 } 1437 1438 // 1439 // backward ordering 1440 // 1441 Last = OrderedCollectionMax (Tree); 1442 BackwardCount = (Last != NULL); 1443 for (Node = OrderedCollectionPrev (Last); Node != NULL; 1444 Node = OrderedCollectionPrev (Last)) { 1445 ASSERT (Tree->UserStructCompare (Last->UserStruct, Node->UserStruct) > 0); 1446 Last = Node; 1447 ++BackwardCount; 1448 } 1449 1450 ASSERT (ForwardCount == BackwardCount); 1451 1452 DEBUG ((DEBUG_VERBOSE, "%a: Tree=%p BlackHeight=%Ld Count=%Ld\n", 1453 __FUNCTION__, Tree, (INT64)BlackHeight, (INT64)ForwardCount)); 1454} 1455