eafbb80fcebb29e5078076a790d834cd4d566ecd
[wimlib] / src / avl_tree.c
1 /*
2  * avl_tree.c
3  *
4  * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
5  * tree), implementation file.
6  *
7  * Author:  Eric Biggers
8  * Year:    2014
9  *
10  * This file is placed into the public domain.  You can do whatever you want
11  * with it.
12  */
13
14 #include <wimlib/avl_tree.h>
15
16 /* Starts an in-order traversal of the tree: returns the least-valued node, or
17  * NULL if the tree is empty.  */
18 struct avl_tree_node *
19 avl_tree_first_in_order(const struct avl_tree_node *root)
20 {
21         const struct avl_tree_node *first = root;
22
23         if (first)
24                 while (first->left)
25                         first = first->left;
26         return (struct avl_tree_node *)first;
27 }
28
29 /* Continues an in-order traversal of the tree: returns the next-greatest-valued
30  * node, or NULL if there is none.  */
31 struct avl_tree_node *
32 avl_tree_next_in_order(const struct avl_tree_node *prev)
33 {
34         const struct avl_tree_node *next;
35
36         if (prev->right)
37                 for (next = prev->right;
38                      next->left;
39                      next = next->left)
40                         ;
41         else
42                 for (next = avl_get_parent(prev);
43                      next && prev == next->right;
44                      prev = next, next = avl_get_parent(next))
45                         ;
46         return (struct avl_tree_node *)next;
47 }
48
49 /* Starts a postorder traversal of the tree.  */
50 struct avl_tree_node *
51 avl_tree_first_in_postorder(const struct avl_tree_node *root)
52 {
53         const struct avl_tree_node *first = root;
54
55         if (first)
56                 while (first->left || first->right)
57                         first = first->left ? first->left : first->right;
58
59         return (struct avl_tree_node *)first;
60 }
61
62 /* Continues a postorder traversal of the tree.  @prev will not be deferenced as
63  * it's allowed that its memory has been freed; @prev_parent must be its saved
64  * parent node.  Returns NULL if there are no more nodes (i.e. @prev was the
65  * root of the tree).  */
66 struct avl_tree_node *
67 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
68                            const struct avl_tree_node *prev_parent)
69 {
70         const struct avl_tree_node *next = prev_parent;
71
72         if (next && prev == next->left && next->right)
73                 for (next = next->right;
74                      next->left || next->right;
75                      next = next->left ? next->left : next->right)
76                         ;
77         return (struct avl_tree_node *)next;
78 }
79
80 /* Get the left child (sign < 0) or the right child (sign > 0)
81  * Note: for all calls of this, 'sign' is constant at compilation time and the
82  * compiler can remove the conditional.  */
83 static AVL_INLINE struct avl_tree_node *
84 avl_get_child(const struct avl_tree_node *parent, int sign)
85 {
86         if (sign < 0)
87                 return parent->left;
88         else
89                 return parent->right;
90 }
91
92 /* Set the left child (sign < 0) or the right child (sign > 0)
93  * Note: for all calls of this, 'sign' is constant at compilation time and the
94  * compiler can remove the conditional.  */
95 static AVL_INLINE void
96 avl_set_child(struct avl_tree_node *parent, int sign,
97               struct avl_tree_node *child)
98 {
99         if (sign < 0)
100                 parent->left = child;
101         else
102                 parent->right = child;
103 }
104
105 /* Sets the parent of the AVL tree @node to @parent.  */
106 static AVL_INLINE void
107 avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
108 {
109         node->parent_balance =
110                 (node->parent_balance & 3) | (uintptr_t)parent;
111 }
112
113 /* Returns the balance factor of the specified AVL tree node --- that is, the
114  * height of its right subtree minus the height of its left subtree.  */
115 static AVL_INLINE int
116 avl_get_balance_factor(const struct avl_tree_node *node)
117 {
118         return (int)(node->parent_balance & 3) - 1;
119 }
120
121 /* Sets the balance factor of the specified AVL tree node.  The caller MUST
122  * ensure that this number is valid and maintains the AVL tree invariants.  */
123 static AVL_INLINE void
124 avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
125 {
126         node->parent_balance =
127                 (node->parent_balance & ~3) | (balance_factor + 1);
128 }
129
130 /* Increments the balance factor of the specified AVL tree @node by the
131  * specified @amount.  The caller MUST ensure that this number is valid and
132  * maintains the AVL tree invariants.  */
133 static AVL_INLINE void
134 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
135 {
136         node->parent_balance += amount;
137 }
138
139 /*
140  * Template for performing a rotation ---
141  *
142  * Clockwise/Right (sign > 0):
143  *
144  *           X             X
145  *           |             |
146  *           A             B
147  *          / \           / \
148  *         B   C   =>    D   A
149  *        / \               / \
150  *       D   E             E   C
151  *
152  * Counterclockwise/Left (sign < 0)- same, except left and right are reversed:
153  *
154  *           X             X
155  *           |             |
156  *           A             B
157  *          / \           / \
158  *         C   B   =>    A   D
159  *            / \       / \
160  *           E   D     C   E
161  *
162  * This updates pointers but not balance factors!
163  */
164 static AVL_INLINE void
165 avl_rotate(struct avl_tree_node * const A, const int sign)
166 {
167         struct avl_tree_node * const B = avl_get_child(A, -sign);
168         struct avl_tree_node * const C = avl_get_child(A, +sign);
169         struct avl_tree_node * const E = avl_get_child(B, +sign);
170         struct avl_tree_node * const X = avl_get_parent(A);
171
172         avl_set_child(A, -sign, E);
173         avl_set_child(A, +sign, C);
174         avl_set_parent(A, B);
175
176         avl_set_child(B, +sign, A);
177         avl_set_parent(B, X);
178
179         if (E)
180                 avl_set_parent(E, A);
181
182         if (X) {
183                 if (avl_get_child(X, +sign) == A)
184                         avl_set_child(X, +sign, B);
185                 else
186                         avl_set_child(X, -sign, B);
187         }
188 }
189
190 /* See description in avl_handle_subtree_growth()  */
191 static AVL_INLINE struct avl_tree_node *
192 avl_do_double_rotate(struct avl_tree_node * const B,
193                      struct avl_tree_node * const A, const int sign)
194 {
195
196         struct avl_tree_node * const E = avl_get_child(B, -sign);
197         const int e = avl_get_balance_factor(E);
198
199         avl_rotate(B, +sign);
200         avl_rotate(A, -sign);
201         avl_set_balance_factor(A, ((sign * e <= 0) ? 0 : -e));
202         avl_set_balance_factor(B, ((sign * e >= 0) ? 0 : -e));
203         avl_set_balance_factor(E, 0);
204
205         return E;
206 }
207
208 static AVL_INLINE bool
209 avl_handle_subtree_growth(struct avl_tree_node * const node,
210                           struct avl_tree_node * const parent,
211                           const int sign)
212 {
213         /* Height of @node subtree of @parent has increased by 1.
214          * Adjust @parent's balance factor and check whether rotations need to
215          * be done.  */
216
217         const int old_balance_factor = avl_get_balance_factor(parent);
218         const int new_balance_factor = old_balance_factor + sign;
219
220         if (old_balance_factor == 0) {
221                 /* Parent has increased in height but is still sufficiently
222                  * balanced.  Continue up the tree.  */
223                 avl_adjust_balance_factor(parent, sign);
224                 return false;
225         }
226
227         if (new_balance_factor == 0) {
228                 /* Parent is balanced; nothing more to to.  */
229                 avl_adjust_balance_factor(parent, sign);
230                 return true;
231         }
232
233         /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
234          * The other case is symmetric --- that is, the rotations done are the
235          * the mirror images, all the balance factors are inverted, and left and
236          * right pointers are otherwise reversed.  */
237
238         /* Parent is left-heavy (balance_factor == -2).  */
239
240         if (sign * avl_get_balance_factor(node) > 0) {
241
242                 /* Child node (@node, also B below) is also left-heavy.
243                  * It must have balance_factor == -1.
244                  * Do a clockwise ("right") rotation rooted at
245                  * @parent (A below):
246                  *
247                  *           A              B
248                  *          / \           /   \
249                  *         B   C   =>    D     A
250                  *        / \           / \   / \
251                  *       D   E         F   G E   C
252                  *      / \
253                  *     F   G
254                  *
255                  * Before the rotation:
256                  *      balance(A) = -2
257                  *      balance(B) = -1
258                  * Let x = height(C).  Then:
259                  *      height(B) = x + 2
260                  *      height(D) = x + 1
261                  *      height(E) = x
262                  *      max(height(F), height(G)) = x.
263                  *
264                  * After the rotation:
265                  *      height(D) = max(height(F), height(G)) + 1
266                  *                = x + 1
267                  *      height(A) = max(height(E), height(C)) + 1
268                  *                = max(x, x) + 1 = x + 1
269                  *      balance(B) = 0
270                  *      balance(A) = 0
271                  */
272
273                 avl_rotate(parent, -sign);
274
275                 avl_set_balance_factor(node, 0);   /* B */
276                 avl_set_balance_factor(parent, 0); /* A */
277         } else {
278                 /* Child node (@node, also B below) is right-heavy.
279                  * It must have balance_factor == +1.
280                  * Do a counterclockwise ("left") rotation rooted at child node
281                  * (B below), then a clockwise ("right") rotation rooted at
282                  * parent node (A below).
283                  *
284                  *           A             A           E
285                  *          / \           / \        /   \
286                  *         B   C   =>    E   C  =>  B     A
287                  *        / \           / \        / \   / \
288                  *       D   E         B   G      D   F G   C
289                  *          / \       / \
290                  *         F   G     D   F
291                  *
292                  * Before the rotation:
293                  *      balance(A) = -2
294                  *      balance(B) = +1
295                  * Let x = height(C).  Then:
296                  *      height(B) = x + 2
297                  *      height(E) = x + 1
298                  *      height(D) = x
299                  *      max(height(F), height(G)) = x
300                  *
301                  * After both rotations:
302                  *      height(A) = max(height(G), height(C)) + 1
303                  *                = x + 1
304                  *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
305                  *      height(B) = max(height(D), height(F)) + 1
306                  *                = x + 1
307                  *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
308                  *
309                  *      height(E) = x + 2
310                  *      balance(E) = 0
311                  */
312                 avl_do_double_rotate(node, parent, sign);
313         }
314         return true;
315 }
316
317 /* Rebalance the tree after insertion of the specified node.  */
318 void
319 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
320                                 struct avl_tree_node *inserted)
321 {
322         struct avl_tree_node *node, *parent;
323         bool done = false;
324
325         inserted->left = NULL;
326         inserted->right = NULL;
327
328         for (node = inserted, parent = avl_get_parent(node);
329              parent && !done;
330              node = parent, parent = avl_get_parent(parent))
331         {
332                 /* Height of @node subtree has increased by 1  */
333
334                 if (node == parent->left)
335                         done = avl_handle_subtree_growth(node, parent, -1);
336                 else
337                         done = avl_handle_subtree_growth(node, parent, +1);
338         }
339         /* Due to rotations, *root_ptr may no longer be the root of the tree  */
340         while (avl_get_parent(*root_ptr))
341                 *root_ptr = avl_get_parent(*root_ptr);
342 }
343
344 static AVL_INLINE struct avl_tree_node *
345 avl_handle_subtree_shrink(struct avl_tree_node *parent,
346                           const int sign,
347                           bool * const left_deleted_ret)
348 {
349         struct avl_tree_node *node;
350
351         const int old_balance_factor = avl_get_balance_factor(parent);
352         const int new_balance_factor = old_balance_factor + sign;
353
354         if (old_balance_factor == 0) {
355                 /* Prior to the deletion, the subtree rooted at
356                  * @parent was perfectly balanced.  It's now
357                  * unbalanced by 1, but that's okay and its height
358                  * hasn't changed.  Nothing more to do.  */
359                 avl_adjust_balance_factor(parent, sign);
360                 return NULL;
361         } else if (new_balance_factor == 0) {
362                 /* The subtree rooted at @parent is now perfectly
363                  * balanced, whereas before the deletion it was
364                  * unbalanced by 1.  Its height must have decreased
365                  * by 1.  No rotation is needed at this location,
366                  * but continue up the tree.  */
367                 avl_adjust_balance_factor(parent, sign);
368                 node = parent;
369         } else {
370                 /* The subtree rooted at @parent is now significantly
371                  * unbalanced (by 2 in some direction).  */
372                 node = avl_get_child(parent, sign);
373
374                 /* The rotations below are similar to those done during
375                  * insertion.  The only new case is the one where the
376                  * child node has a balance factor of 0.  */
377
378                 if (sign * avl_get_balance_factor(node) >= 0) {
379                         avl_rotate(parent, -sign);
380
381                         if (avl_get_balance_factor(node) == 0) {
382                                 avl_set_balance_factor(node,   -sign);
383                                 avl_set_balance_factor(parent, +sign);
384                                 /* Height is unchanged; nothing more to do.  */
385                                 return NULL;
386                         } else {
387                                 avl_set_balance_factor(node, 0);
388                                 avl_set_balance_factor(parent, 0);
389                         }
390                 } else {
391                         node = avl_do_double_rotate(node, parent, sign);
392                 }
393         }
394         parent = avl_get_parent(node);
395         if (parent)
396                 *left_deleted_ret = (node == parent->left);
397         return parent;
398 }
399
400 /* Swaps node X, which must have 2 children, with its in-order successor.
401  * This adjusts all the necessary pointers and balance factors.  */
402 static AVL_INLINE void
403 avl_tree_swap_with_successor(struct avl_tree_node * const X)
404 {
405         struct avl_tree_node * const Y          = avl_tree_next_in_order(X);
406         struct avl_tree_node * const Y_right    = Y->right;
407         struct avl_tree_node * const Y_parent   = avl_get_parent(Y);
408         struct avl_tree_node * const X_parent   = avl_get_parent(X);
409         const int Y_balance_factor              = avl_get_balance_factor(Y);
410
411         /* Set child pointer to Y when placed in X's position  */
412         if (X_parent) {
413                 if (X == X_parent->left)
414                         X_parent->left = Y;
415                 else
416                         X_parent->right = Y;
417         }
418
419         /* Set child pointer to X when placed in Y's position.  Also set
420          * the right pointer of Y (it may point to X if the nodes are
421          * adjacent)  */
422         if (Y_parent != X) {
423                 if (Y == Y_parent->left)
424                         Y_parent->left = X;
425                 else
426                         Y_parent->right = X;
427                 Y->right = X->right;
428         } else {
429                 Y->right = X;
430         }
431
432         /* Set parent pointer to X when placed in Y's position  */
433         if (Y_right)
434                 avl_set_parent(Y_right, X);
435
436         /* Note: Y has no left child since it is the in-order sucessor of X.  */
437
438         /* Set parent pointers to Y when placed in X's position, as well as X's
439          * parent when placed in Y's position.  */
440         avl_set_parent(X->left, Y);
441         if (X->right != Y) {
442                 avl_set_parent(X->right, Y);
443                 avl_set_parent(X, Y_parent);
444         } else {
445                 avl_set_parent(X, Y);
446         }
447
448         avl_set_parent(Y, X_parent);
449         Y->left = X->left;
450         avl_set_balance_factor(Y, avl_get_balance_factor(X));
451
452         X->left = NULL;
453         X->right = Y_right;
454         avl_set_balance_factor(X, Y_balance_factor);
455 }
456
457 /* Removes the specified @node from the AVL tree.  @root_ptr must point to the
458  * pointer to the root node of the tree; *root_ptr may change if the tree is
459  * rebalanced.
460  *
461  * This *only* removes the node and rebalances the tree; it does not free
462  * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
463 void
464 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
465 {
466         struct avl_tree_node *child, *parent;
467         bool left_deleted;
468
469         if (node->left && node->right)
470                 avl_tree_swap_with_successor(node);
471
472         /* Unlink @node  */
473         child = node->left ? node->left : node->right;
474         parent = avl_get_parent(node);
475         if (parent) {
476                 if (node == parent->left) {
477                         parent->left = child;
478                         left_deleted = true;
479                 } else {
480                         parent->right = child;
481                         left_deleted = false;
482                 }
483         } else {
484                 *root_ptr = child;
485         }
486         if (child)
487                 avl_set_parent(child, parent);
488
489         /* Rebalance the tree  */
490         while (parent) {
491                 if (left_deleted)
492                         parent = avl_handle_subtree_shrink(parent, +1, &left_deleted);
493                 else
494                         parent = avl_handle_subtree_shrink(parent, -1, &left_deleted);
495         }
496
497         /* Due to rotations, *root_ptr may no longer point to the root of the
498          * tree.  Fix it.  */
499         if (*root_ptr)
500                 while (avl_get_parent(*root_ptr))
501                         *root_ptr = avl_get_parent(*root_ptr);
502 }