xref: /reactos/sdk/lib/rtl/avlsupp.c (revision 40462c92)
1 /*
2  * PROJECT:         ReactOS Runtime Library
3  * LICENSE:         BSD - See COPYING.ARM in the top level directory
4  * FILE:            lib/rtl/avlsupp.c
5  * PURPOSE:         AVL Tree Internal Support Routines/Main Algorithms
6  * PROGRAMMERS:     ReactOS Portable Systems Group
7  */
8 
9 /* INCLUDES ******************************************************************/
10 
11 /* Internal header for table entries */
12 typedef struct _TABLE_ENTRY_HEADER
13 {
14     RTL_BALANCED_LINKS BalancedLinks;
15     LONGLONG UserData;
16 } TABLE_ENTRY_HEADER, *PTABLE_ENTRY_HEADER;
17 
18 typedef enum _RTL_AVL_BALANCE_FACTOR
19 {
20     RtlUnbalancedAvlTree = -2,
21     RtlLeftHeavyAvlTree,
22     RtlBalancedAvlTree,
23     RtlRightHeavyAvlTree,
24 } RTL_AVL_BALANCE_FACTOR;
25 
26 C_ASSERT(RtlBalancedAvlTree == 0);
27 
28 /* FUNCTIONS ******************************************************************/
29 
30 FORCEINLINE
31 TABLE_SEARCH_RESULT
32 RtlpFindAvlTableNodeOrParent(IN PRTL_AVL_TABLE Table,
33                              IN PVOID Buffer,
34                              OUT PRTL_BALANCED_LINKS *NodeOrParent)
35 {
36     PRTL_BALANCED_LINKS CurrentNode, ChildNode;
37     RTL_GENERIC_COMPARE_RESULTS Result;
38 
39     /* Quick check to see if the table is empty */
40     if (!Table->NumberGenericTableElements) return TableEmptyTree;
41 
42     /* Set the current node */
43     CurrentNode = RtlRightChildAvl(&Table->BalancedRoot);
44 
45     /* Start compare loop */
46     while (TRUE)
47     {
48         /* Compare which side is greater */
49         Result = RtlpAvlCompareRoutine(Table,
50                                        Buffer,
51                                        &((PTABLE_ENTRY_HEADER)CurrentNode)->
52                                        UserData);
53         if (Result == GenericLessThan)
54         {
55             /* We're less, check if this is the left child */
56             ChildNode = RtlLeftChildAvl(CurrentNode);
57             if (ChildNode)
58             {
59                 /* Continue searching from this node */
60                 CurrentNode = ChildNode;
61             }
62             else
63             {
64                 /* Otherwise, the element isn't in this tree */
65                 *NodeOrParent = CurrentNode;
66                 return TableInsertAsLeft;
67             }
68         }
69         else if (Result == GenericGreaterThan)
70         {
71             /* We're more, check if this is the right child */
72             ChildNode = RtlRightChildAvl(CurrentNode);
73             if (ChildNode)
74             {
75                 /* Continue searching from this node */
76                 CurrentNode = ChildNode;
77             }
78             else
79             {
80                 /* Otherwise, the element isn't in this tree */
81                 *NodeOrParent = CurrentNode;
82                 return TableInsertAsRight;
83             }
84         }
85         else
86         {
87             /* We should've found the node */
88             ASSERT(Result == GenericEqual);
89 
90             /* Return node found */
91             *NodeOrParent = CurrentNode;
92             return TableFoundNode;
93         }
94     }
95 }
96 
97 FORCEINLINE
98 VOID
99 RtlpPromoteAvlTreeNode(IN PRTL_BALANCED_LINKS Node)
100 {
101     PRTL_BALANCED_LINKS ParentNode, SuperParentNode;
102     PRTL_BALANCED_LINKS *SwapNode1, *SwapNode2;
103 
104     /* Grab parents up to 2 levels high */
105     ParentNode = RtlParentAvl(Node);
106     SuperParentNode = RtlParentAvl(ParentNode);
107 
108     /* Pick which nodes will be rotated */
109     SwapNode1 = RtlIsLeftChildAvl(Node) ? &ParentNode->LeftChild : &ParentNode->RightChild;
110     SwapNode2 = RtlIsLeftChildAvl(Node) ? &Node->RightChild : &Node->LeftChild;
111 
112     /* Do the rotate, and update the parent and super-parent as needed */
113     *SwapNode1 = *SwapNode2;
114     if (*SwapNode1) RtlSetParent(*SwapNode1, ParentNode);
115     *SwapNode2 = ParentNode;
116     RtlSetParent(ParentNode, Node);
117 
118     /* Now update the super-parent child link, and make it parent of the node*/
119     SwapNode1 = (RtlLeftChildAvl(SuperParentNode) == ParentNode) ?
120                  &SuperParentNode->LeftChild: &SuperParentNode->RightChild;
121     *SwapNode1 = Node;
122     RtlSetParent(Node, SuperParentNode);
123 }
124 
125 FORCEINLINE
126 BOOLEAN
127 RtlpRebalanceAvlTreeNode(IN PRTL_BALANCED_LINKS Node)
128 {
129     PRTL_BALANCED_LINKS ChildNode, SubChildNode;
130     CHAR Balance;
131     ASSERT(RtlParentAvl(Node) != Node);
132 
133     /* Get the balance, and figure out which child node to go down on */
134     Balance = RtlBalance(Node);
135     ChildNode = (Balance == RtlRightHeavyAvlTree) ?
136                  RtlRightChildAvl(Node) : RtlLeftChildAvl(Node);
137 
138     /* The child and node have the same balance, promote the child upwards */
139     if (RtlBalance(ChildNode) == Balance)
140     {
141         /* This performs the rotation described in Knuth A8-A10 for Case 1 */
142         RtlpPromoteAvlTreeNode(ChildNode);
143 
144         /* The nodes are now balanced */
145         RtlSetBalance(ChildNode, RtlBalancedAvlTree);
146         RtlSetBalance(Node, RtlBalancedAvlTree);
147         return FALSE;
148     }
149 
150     /* The child has the opposite balance, a double promotion of the child's child must happen */
151     if (RtlBalance(ChildNode) == -Balance)
152     {
153         /* Pick which sub-child to use based on the balance */
154         SubChildNode = (Balance == RtlRightHeavyAvlTree) ?
155                         RtlLeftChildAvl(ChildNode) : RtlRightChildAvl(ChildNode);
156 
157         /* Do the double-rotation described in Knuth A8-A10 for Case 2 */
158         RtlpPromoteAvlTreeNode(SubChildNode);
159         RtlpPromoteAvlTreeNode(SubChildNode);
160 
161         /* Was the sub-child sharing the same balance as the node? */
162         if (RtlBalance(SubChildNode) == Balance)
163         {
164             /* Then the subchild is now balanced, and the node's weight is inversed */
165             RtlSetBalance(ChildNode, RtlBalancedAvlTree);
166             RtlSetBalance(Node, -Balance);
167         }
168         else if (RtlBalance(SubChildNode) == -Balance)
169         {
170             /*
171              * In this case, the sub-child weight was the inverse of the node, so
172              * the child now shares the node's balance original weight, while the
173              * node becomes balanced.
174              */
175             RtlSetBalance(ChildNode, Balance);
176             RtlSetBalance(Node, RtlBalancedAvlTree);
177         }
178         else
179         {
180             /*
181              * Otherwise, the sub-child was unbalanced, so both the child and node
182              * now become balanced.
183              */
184             RtlSetBalance(ChildNode, RtlBalancedAvlTree);
185             RtlSetBalance(Node, RtlBalancedAvlTree);
186         }
187 
188         /* In all cases, the sub-child is now balanced */
189         RtlSetBalance(SubChildNode, RtlBalancedAvlTree);
190         return FALSE;
191     }
192 
193     /*
194      * The case that remains is that the child was already balanced, so this is
195      * This is the rotation required for Case 3 in Knuth A8-A10
196      */
197     RtlpPromoteAvlTreeNode(ChildNode);
198 
199     /* Now the child has the opposite weight of the node */
200     RtlSetBalance(ChildNode, -Balance);
201 
202     /* This only happens on deletion, so we return TRUE to terminate the delete */
203     return TRUE;
204 }
205 
206 FORCEINLINE
207 VOID
208 RtlpInsertAvlTreeNode(IN PRTL_AVL_TABLE Table,
209                       IN PRTL_BALANCED_LINKS NewNode,
210                       IN OUT PVOID NodeOrParent,
211                       IN OUT TABLE_SEARCH_RESULT SearchResult)
212 {
213     CHAR Balance;
214 
215     /* Initialize the new inserted element */
216     MI_ASSERT(SearchResult != TableFoundNode);
217     NewNode->LeftChild = NewNode->RightChild = NULL;
218     RtlSetBalance(NewNode, RtlBalancedAvlTree);
219 
220     /* Increase element count */
221     Table->NumberGenericTableElements++;
222 
223     /* Check where we should insert the entry */
224     if (SearchResult == TableEmptyTree)
225     {
226         /* This is the new root node */
227         RtlInsertAsRightChildAvl(&Table->BalancedRoot, NewNode);
228 
229         /* On AVL trees, we also update the depth */
230         ASSERT(Table->DepthOfTree == 0);
231         Table->DepthOfTree = 1;
232         return;
233     }
234     else if (SearchResult == TableInsertAsLeft)
235     {
236         /* Insert it left */
237         RtlInsertAsLeftChildAvl(NodeOrParent, NewNode);
238     }
239     else
240     {
241         /* Right node */
242         RtlInsertAsRightChildAvl(NodeOrParent, NewNode);
243     }
244 
245     /* Little cheat to save on loop processing, taken from Timo */
246     RtlSetBalance(&Table->BalancedRoot, RtlLeftHeavyAvlTree);
247 
248     /*
249      * This implements A6-A7 from Knuth based on http://coding.derkeiler.com
250      * /pdf/Archive/C_CPP/comp.lang.c/2004-01/1812.pdf, however the algorithm
251      * is slightly modified to follow the tree based on the Parent Node such
252      * as the Windows algorithm does it, instead of following the nodes down.
253      */
254     while (TRUE)
255     {
256         /* Calculate which side to balance on */
257         Balance = RtlIsLeftChildAvl(NewNode) ? RtlLeftHeavyAvlTree : RtlRightHeavyAvlTree;
258 
259         /* Check if the parent node was balanced */
260         if (RtlBalance(NodeOrParent) == RtlBalancedAvlTree)
261         {
262             /* It's not balanced anymore (heavy on one side) */
263             RtlSetBalance(NodeOrParent, Balance);
264 
265             /* Move up */
266             NewNode = NodeOrParent;
267             NodeOrParent = RtlParentAvl(NodeOrParent);
268         }
269         else if (RtlBalance(NodeOrParent) != Balance)
270         {
271             /* The parent's balance is opposite, so the tree is balanced now */
272             RtlSetBalance(NodeOrParent, RtlBalancedAvlTree);
273 
274             /* Check if this is the root (the cheat applied earlier gets us here) */
275             if (RtlBalance(&Table->BalancedRoot) == RtlBalancedAvlTree)
276             {
277                 /* The depth has thus increased */
278                 Table->DepthOfTree++;
279             }
280 
281             /* We reached the root or a balanced node, so we're done */
282             break;
283         }
284         else
285         {
286             /* The tree is now unbalanced, so AVL rebalancing must happen */
287             RtlpRebalanceAvlTreeNode(NodeOrParent);
288             break;
289         }
290     }
291 }
292 
293 FORCEINLINE
294 VOID
295 RtlpDeleteAvlTreeNode(IN PRTL_AVL_TABLE Table,
296                       IN PRTL_BALANCED_LINKS Node)
297 {
298     PRTL_BALANCED_LINKS DeleteNode = NULL, ParentNode;
299     PRTL_BALANCED_LINKS *Node1, *Node2;
300     CHAR Balance;
301 
302     /* Take one of the children if possible */
303     if (!(RtlLeftChildAvl(Node)) || !(RtlRightChildAvl(Node))) DeleteNode = Node;
304 
305     /* Otherwise, check if one side is longer */
306     if (!(DeleteNode) && (RtlBalance(Node) >= RtlBalancedAvlTree))
307     {
308         /* Pick the successor which will be the longest side in this case */
309         DeleteNode = RtlRightChildAvl(Node);
310         while (RtlLeftChildAvl(DeleteNode)) DeleteNode = RtlLeftChildAvl(DeleteNode);
311     }
312     else if (!DeleteNode)
313     {
314         /* Pick the predecessor which will be the longest side in this case */
315         DeleteNode = RtlLeftChildAvl(Node);
316         while (RtlRightChildAvl(DeleteNode)) DeleteNode = RtlRightChildAvl(DeleteNode);
317     }
318 
319     /* Get the parent node */
320     ParentNode = RtlParentAvl(DeleteNode);
321     DPRINT("Parent: %p\n", ParentNode);
322 
323     /* Pick which now to use based on whether or not we have a left child */
324     Node1 = RtlLeftChildAvl(DeleteNode) ? &DeleteNode->LeftChild : &DeleteNode->RightChild;
325     DPRINT("Node 1: %p %p\n", Node1, *Node1);
326 
327     /* Pick which node to swap based on if we're already a left child or not */
328     Node2 = RtlIsLeftChildAvl(DeleteNode) ? &ParentNode->LeftChild : &ParentNode->RightChild;
329     DPRINT("Node 2: %p %p\n", Node2, *Node2);
330 
331     /* Pick the correct balance depending on which side will get heavier */
332     Balance = RtlIsLeftChildAvl(DeleteNode) ? RtlLeftHeavyAvlTree : RtlRightHeavyAvlTree;
333     DPRINT("Balance: %lx\n", Balance);
334 
335     /* Swap the children nodes, making one side heavier */
336     *Node2 = *Node1;
337 
338     /* If the node has a child now, update its parent */
339     if (*Node1) RtlSetParent(*Node1, ParentNode);
340 
341     /* Assume balanced root for loop optimization */
342     RtlSetBalance(&Table->BalancedRoot, RtlBalancedAvlTree);
343 
344     /* Loop up the tree by parents */
345     while (TRUE)
346     {
347         /* Check if the tree's balance increased */
348         if (RtlBalance(ParentNode) == Balance)
349         {
350             /* Now the tree is balanced */
351             RtlSetBalance(ParentNode, RtlBalancedAvlTree);
352         }
353         else if (RtlBalance(ParentNode) == RtlBalancedAvlTree)
354         {
355             /* The tree has now become less balanced, since it was balanced */
356             RtlSetBalance(ParentNode, -Balance);
357 
358             /* Deal with the loop optimization to detect loss of a tree level */
359             if (RtlBalance(&Table->BalancedRoot) != RtlBalancedAvlTree) Table->DepthOfTree--;
360             break;
361         }
362         else
363         {
364             /* The tree has become unbalanced, so a rebalance is needed */
365             if (RtlpRebalanceAvlTreeNode(ParentNode)) break;
366 
367             /* Get the new parent after the balance */
368             ParentNode = RtlParentAvl(ParentNode);
369         }
370 
371         /* Choose which balance factor to use based on which side we're on */
372         Balance = RtlIsRightChild(ParentNode) ?
373                   RtlRightHeavyAvlTree : RtlLeftHeavyAvlTree;
374 
375         /* Iterate up the tree */
376         ParentNode = RtlParentAvl(ParentNode);
377     }
378 
379     /* Check if this isn't the node we ended up deleting directly */
380     if (Node == DeleteNode) return;
381 
382     /* Copy the deleted node itself */
383     RtlpCopyAvlNodeData(DeleteNode, Node);
384 
385     /* Pick the right node to unlink */
386     Node1 = RtlIsLeftChildAvl(Node) ?
387             &(RtlParentAvl(DeleteNode))->LeftChild : &(RtlParentAvl(DeleteNode))->RightChild;
388     *Node1 = DeleteNode;
389 
390     /* Reparent as appropriate */
391     if (RtlLeftChildAvl(DeleteNode)) RtlSetParent(RtlLeftChildAvl(DeleteNode), DeleteNode);
392     if (RtlRightChildAvl(DeleteNode)) RtlSetParent(RtlRightChildAvl(DeleteNode), DeleteNode);
393 }
394 
395 /* EOF */
396