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
29 typedef 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 //
39 typedef ORDERED_COLLECTION RED_BLACK_TREE;
40 typedef ORDERED_COLLECTION_ENTRY RED_BLACK_TREE_NODE;
41 typedef ORDERED_COLLECTION_USER_COMPARE RED_BLACK_TREE_USER_COMPARE;
42 typedef ORDERED_COLLECTION_KEY_COMPARE RED_BLACK_TREE_KEY_COMPARE;
43
44 struct ORDERED_COLLECTION {
45 RED_BLACK_TREE_NODE *Root;
46 RED_BLACK_TREE_USER_COMPARE UserStructCompare;
47 RED_BLACK_TREE_KEY_COMPARE KeyCompare;
48 };
49
50 struct 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 **/
70 VOID *
71 EFIAPI
OrderedCollectionUserStruct(IN CONST RED_BLACK_TREE_NODE * Node)72 OrderedCollectionUserStruct (
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 **/
90 VOID
91 RedBlackTreeValidate (
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 **/
115 RED_BLACK_TREE *
116 EFIAPI
OrderedCollectionInit(IN RED_BLACK_TREE_USER_COMPARE UserStructCompare,IN RED_BLACK_TREE_KEY_COMPARE KeyCompare)117 OrderedCollectionInit (
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 **/
151 BOOLEAN
152 EFIAPI
OrderedCollectionIsEmpty(IN CONST RED_BLACK_TREE * Tree)153 OrderedCollectionIsEmpty (
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 **/
173 VOID
174 EFIAPI
OrderedCollectionUninit(IN RED_BLACK_TREE * Tree)175 OrderedCollectionUninit (
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 **/
201 RED_BLACK_TREE_NODE *
202 EFIAPI
OrderedCollectionFind(IN CONST RED_BLACK_TREE * Tree,IN CONST VOID * StandaloneKey)203 OrderedCollectionFind (
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 **/
237 RED_BLACK_TREE_NODE *
238 EFIAPI
OrderedCollectionMin(IN CONST RED_BLACK_TREE * Tree)239 OrderedCollectionMin (
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 **/
269 RED_BLACK_TREE_NODE *
270 EFIAPI
OrderedCollectionMax(IN CONST RED_BLACK_TREE * Tree)271 OrderedCollectionMax (
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 **/
302 RED_BLACK_TREE_NODE *
303 EFIAPI
OrderedCollectionNext(IN CONST RED_BLACK_TREE_NODE * Node)304 OrderedCollectionNext (
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 **/
355 RED_BLACK_TREE_NODE *
356 EFIAPI
OrderedCollectionPrev(IN CONST RED_BLACK_TREE_NODE * Node)357 OrderedCollectionPrev (
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 **/
426 VOID
RedBlackTreeRotateRight(IN OUT RED_BLACK_TREE_NODE * Pivot,OUT RED_BLACK_TREE_NODE ** NewRoot)427 RedBlackTreeRotateRight (
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 **/
491 VOID
RedBlackTreeRotateLeft(IN OUT RED_BLACK_TREE_NODE * Pivot,OUT RED_BLACK_TREE_NODE ** NewRoot)492 RedBlackTreeRotateLeft (
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 **/
585 RETURN_STATUS
586 EFIAPI
OrderedCollectionInsert(IN OUT RED_BLACK_TREE * Tree,OUT RED_BLACK_TREE_NODE ** Node OPTIONAL,IN VOID * UserStruct)587 OrderedCollectionInsert (
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
827 Done:
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 **/
844 BOOLEAN
NodeIsNullOrBlack(IN CONST RED_BLACK_TREE_NODE * Node)845 NodeIsNullOrBlack (
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 **/
918 VOID
919 EFIAPI
OrderedCollectionDelete(IN OUT RED_BLACK_TREE * Tree,IN RED_BLACK_TREE_NODE * Node,OUT VOID ** UserStruct OPTIONAL)920 OrderedCollectionDelete (
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 **/
1353 UINT32
RedBlackTreeRecursiveCheck(IN CONST RED_BLACK_TREE_NODE * Node)1354 RedBlackTreeRecursiveCheck (
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 **/
1403 VOID
RedBlackTreeValidate(IN CONST RED_BLACK_TREE * Tree)1404 RedBlackTreeValidate (
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