8b8e07bbe5d608b26975a47d60e443c044c5d033
[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 /* Sets the parent and balance factor of the AVL tree @node.  */
114 static AVL_INLINE void
115 avl_set_parent_balance(struct avl_tree_node *node, struct avl_tree_node *parent,
116                        int balance_factor)
117 {
118         node->parent_balance = (uintptr_t)parent | (balance_factor + 1);
119 }
120
121 /* Returns the balance factor of the specified AVL tree node --- that is, the
122  * height of its right subtree minus the height of its left subtree.  */
123 static AVL_INLINE int
124 avl_get_balance_factor(const struct avl_tree_node *node)
125 {
126         return (int)(node->parent_balance & 3) - 1;
127 }
128
129 /* Sets the balance factor of the specified AVL tree node.  The caller MUST
130  * ensure that this number is valid and maintains the AVL tree invariants.  */
131 static AVL_INLINE void
132 avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
133 {
134         node->parent_balance =
135                 (node->parent_balance & ~3) | (balance_factor + 1);
136 }
137
138 /* Increments the balance factor of the specified AVL tree @node by the
139  * specified @amount.  The caller MUST ensure that this number is valid and
140  * maintains the AVL tree invariants.  */
141 static AVL_INLINE void
142 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
143 {
144         node->parent_balance += amount;
145 }
146
147 /*
148  * Template for performing a rotation ---
149  *
150  * Clockwise/Right (sign > 0):
151  *
152  *           X             X
153  *           |             |
154  *           A             B
155  *          / \           / \
156  *         B   C   =>    D   A
157  *        / \               / \
158  *       D   E             E   C
159  *
160  * Counterclockwise/Left (sign < 0)- same, except left and right are reversed:
161  *
162  *           X             X
163  *           |             |
164  *           A             B
165  *          / \           / \
166  *         C   B   =>    A   D
167  *            / \       / \
168  *           E   D     C   E
169  *
170  * This updates pointers but not balance factors!
171  */
172 static AVL_INLINE void
173 avl_rotate(struct avl_tree_node ** const root_ptr,
174            struct avl_tree_node * const A, const int sign)
175 {
176         struct avl_tree_node * const B = avl_get_child(A, -sign);
177         struct avl_tree_node * const C = avl_get_child(A, +sign);
178         struct avl_tree_node * const E = avl_get_child(B, +sign);
179         struct avl_tree_node * const X = avl_get_parent(A);
180
181         avl_set_child(A, -sign, E);
182         avl_set_child(A, +sign, C);
183         avl_set_parent(A, B);
184
185         avl_set_child(B, +sign, A);
186         avl_set_parent(B, X);
187
188         if (E)
189                 avl_set_parent(E, A);
190
191         if (X) {
192                 if (avl_get_child(X, +sign) == A)
193                         avl_set_child(X, +sign, B);
194                 else
195                         avl_set_child(X, -sign, B);
196         } else {
197                 *root_ptr = B;
198         }
199 }
200
201 /*
202  * Template for performing a double rotation ---
203  *
204  * Rotate counterclockwise (left) rooted at B, then clockwise (right) rooted at
205  * A (sign > 0):
206  *
207  *           A             A           E
208  *          / \           / \        /   \
209  *         B   C   =>    E   C  =>  B     A
210  *        / \           / \        / \   / \
211  *       D   E         B   G      D   F G   C
212  *          / \       / \
213  *         F   G     D   F
214  *
215  * Rotate clockwise (right) rooted at B, then counterclockwise (left) rooted at
216  * A (sign < 0):
217  *
218  *         A           A               E
219  *        / \         / \            /   \
220  *       C   B   =>  C   E    =>    A     B
221  *          / \         / \        / \   / \
222  *         E   D       G   B      C   G F   D
223  *        / \             / \
224  *       G   F           F   D
225  *
226  * Returns a pointer to E, the new root, and updates balance factors.  Except
227  * for that, this is equivalent to:
228  *      avl_rotate(root_ptr, B, -sign);
229  *      avl_rotate(root_ptr, A, +sign);
230  *
231  */
232 static AVL_INLINE struct avl_tree_node *
233 avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
234                      struct avl_tree_node * const B,
235                      struct avl_tree_node * const A, const int sign)
236 {
237         struct avl_tree_node * const E = avl_get_child(B, +sign);
238         struct avl_tree_node * const F = avl_get_child(E, -sign);
239         struct avl_tree_node * const G = avl_get_child(E, +sign);
240         struct avl_tree_node * const X = avl_get_parent(A);
241         const int e = avl_get_balance_factor(E);
242
243         avl_set_child(A, -sign, G);
244         avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
245
246         avl_set_child(B, +sign, F);
247         avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
248
249         avl_set_child(E, +sign, A);
250         avl_set_child(E, -sign, B);
251         avl_set_parent_balance(E, X, 0);
252
253         if (G)
254                 avl_set_parent(G, A);
255
256         if (F)
257                 avl_set_parent(F, B);
258
259         if (X) {
260                 if (avl_get_child(X, +sign) == A)
261                         avl_set_child(X, +sign, E);
262                 else
263                         avl_set_child(X, -sign, E);
264         } else {
265                 *root_ptr = E;
266         }
267
268         return E;
269 }
270
271 static AVL_INLINE bool
272 avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
273                           struct avl_tree_node * const node,
274                           struct avl_tree_node * const parent,
275                           const int sign)
276 {
277         /* The subtree rooted at @node, which is a child of @parent, has
278          * increased by 1.
279          *
280          * Adjust @parent's balance factor and check whether rotations need to
281          * be done.  */
282
283         int old_balance_factor, new_balance_factor;
284
285         old_balance_factor = avl_get_balance_factor(parent);
286
287         if (old_balance_factor == 0) {
288                 /* Parent has increased in height but is still sufficiently
289                  * balanced.  Continue up the tree.  */
290                 avl_adjust_balance_factor(parent, sign);
291                 return false;
292         }
293
294         new_balance_factor = old_balance_factor + sign;
295
296         if (new_balance_factor == 0) {
297                 /* Parent is balanced; nothing more to to.  */
298                 avl_adjust_balance_factor(parent, sign);
299                 return true;
300         }
301
302         /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
303          * The other case is symmetric --- that is, the rotations done are the
304          * the mirror images, all the balance factors are inverted, and left and
305          * right pointers are otherwise reversed.  */
306
307         /* Parent is left-heavy (balance_factor == -2).  */
308
309         if (sign * avl_get_balance_factor(node) > 0) {
310
311                 /* Child node (@node, also B below) is also left-heavy.
312                  * It must have balance_factor == -1.
313                  * Do a clockwise ("right") rotation rooted at
314                  * @parent (A below):
315                  *
316                  *           A              B
317                  *          / \           /   \
318                  *         B   C   =>    D     A
319                  *        / \           / \   / \
320                  *       D   E         F   G E   C
321                  *      / \
322                  *     F   G
323                  *
324                  * Before the rotation:
325                  *      balance(A) = -2
326                  *      balance(B) = -1
327                  * Let x = height(C).  Then:
328                  *      height(B) = x + 2
329                  *      height(D) = x + 1
330                  *      height(E) = x
331                  *      max(height(F), height(G)) = x.
332                  *
333                  * After the rotation:
334                  *      height(D) = max(height(F), height(G)) + 1
335                  *                = x + 1
336                  *      height(A) = max(height(E), height(C)) + 1
337                  *                = max(x, x) + 1 = x + 1
338                  *      balance(B) = 0
339                  *      balance(A) = 0
340                  */
341
342                 avl_rotate(root_ptr, parent, -sign);
343                 avl_adjust_balance_factor(parent, -sign); /* A */
344                 avl_adjust_balance_factor(node, -sign);   /* B */
345         } else {
346                 /* Child node (@node, also B below) is right-heavy.
347                  * It must have balance_factor == +1.
348                  * Do a counterclockwise ("left") rotation rooted at child node
349                  * (B below), then a clockwise ("right") rotation rooted at
350                  * parent node (A below).
351                  *
352                  *           A             A           E
353                  *          / \           / \        /   \
354                  *         B   C   =>    E   C  =>  B     A
355                  *        / \           / \        / \   / \
356                  *       D   E         B   G      D   F G   C
357                  *          / \       / \
358                  *         F   G     D   F
359                  *
360                  * Before the rotation:
361                  *      balance(A) = -2
362                  *      balance(B) = +1
363                  * Let x = height(C).  Then:
364                  *      height(B) = x + 2
365                  *      height(E) = x + 1
366                  *      height(D) = x
367                  *      max(height(F), height(G)) = x
368                  *
369                  * After both rotations:
370                  *      height(A) = max(height(G), height(C)) + 1
371                  *                = x + 1
372                  *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
373                  *      height(B) = max(height(D), height(F)) + 1
374                  *                = x + 1
375                  *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
376                  *
377                  *      height(E) = x + 2
378                  *      balance(E) = 0
379                  */
380                 avl_do_double_rotate(root_ptr, node, parent, -sign);
381         }
382         return true;
383 }
384
385 /* Rebalance the tree after insertion of the specified node.  */
386 void
387 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
388                                 struct avl_tree_node *inserted)
389 {
390         struct avl_tree_node *node, *parent;
391         bool done;
392
393         inserted->left = NULL;
394         inserted->right = NULL;
395
396         node = inserted;
397
398         /* Adjust balance factor of new node's parent.
399          * No rotation will need to be done at this level.  */
400
401         parent = avl_get_parent(node);
402         if (!parent)
403                 return;
404
405         if (node == parent->left)
406                 avl_adjust_balance_factor(parent, -1);
407         else
408                 avl_adjust_balance_factor(parent, +1);
409
410         if (avl_get_balance_factor(parent) == 0)
411                 /* @parent did not change in height.  Nothing more to do.  */
412                 return;
413
414         /* The subtree rooted at @parent increased in height by 1.  */
415
416         do {
417                 /* Adjust balance factor of next ancestor.  */
418
419                 node = parent;
420                 parent = avl_get_parent(node);
421                 if (!parent)
422                         return;
423
424                 /* The subtree rooted at @node has increased in height by 1.  */
425                 if (node == parent->left)
426                         done = avl_handle_subtree_growth(root_ptr, node,
427                                                          parent, -1);
428                 else
429                         done = avl_handle_subtree_growth(root_ptr, node,
430                                                          parent, +1);
431         } while (!done);
432 }
433
434 static AVL_INLINE struct avl_tree_node *
435 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
436                           struct avl_tree_node *parent,
437                           const int sign,
438                           bool * const left_deleted_ret)
439 {
440         struct avl_tree_node *node;
441
442         const int old_balance_factor = avl_get_balance_factor(parent);
443         const int new_balance_factor = old_balance_factor + sign;
444
445         if (old_balance_factor == 0) {
446                 /* Prior to the deletion, the subtree rooted at
447                  * @parent was perfectly balanced.  It's now
448                  * unbalanced by 1, but that's okay and its height
449                  * hasn't changed.  Nothing more to do.  */
450                 avl_adjust_balance_factor(parent, sign);
451                 return NULL;
452         } else if (new_balance_factor == 0) {
453                 /* The subtree rooted at @parent is now perfectly
454                  * balanced, whereas before the deletion it was
455                  * unbalanced by 1.  Its height must have decreased
456                  * by 1.  No rotation is needed at this location,
457                  * but continue up the tree.  */
458                 avl_adjust_balance_factor(parent, sign);
459                 node = parent;
460         } else {
461                 /* The subtree rooted at @parent is now significantly
462                  * unbalanced (by 2 in some direction).  */
463                 node = avl_get_child(parent, sign);
464
465                 /* The rotations below are similar to those done during
466                  * insertion.  The only new case is the one where the
467                  * child node has a balance factor of 0.  */
468
469                 if (sign * avl_get_balance_factor(node) >= 0) {
470                         avl_rotate(root_ptr, parent, -sign);
471
472                         if (avl_get_balance_factor(node) == 0) {
473                                 /* Child node (@node, also B below) is balanced.
474                                  * Do a clockwise ("right") rotation rooted at
475                                  * @parent (A below):
476                                  *
477                                  *           A              B
478                                  *          / \           /   \
479                                  *         B   C   =>    D     A
480                                  *        / \           / \   / \
481                                  *       D   E         F   G E   C
482                                  *      / \
483                                  *     F   G
484                                  *
485                                  * Before the rotation:
486                                  *      balance(A) = -2
487                                  *      balance(B) =  0
488                                  * Let x = height(C).  Then:
489                                  *      height(B) = x + 2
490                                  *      height(D) = x + 1
491                                  *      height(E) = x + 1
492                                  *      max(height(F), height(G)) = x.
493                                  *
494                                  * After the rotation:
495                                  *      height(D) = max(height(F), height(G)) + 1
496                                  *                = x + 1
497                                  *      height(A) = max(height(E), height(C)) + 1
498                                  *                = max(x + 1, x) + 1 = x + 2
499                                  *      balance(A) = -1
500                                  *      balance(B) = +1
501                                  */
502
503                                 /* A: -2 => -1 (sign < 0)
504                                  * or +2 => +1 (sign > 0)
505                                  * No change needed --- that's the same as
506                                  * old_balance_factor.  */
507
508                                 /* B: 0 => +1 (sign < 0)
509                                  * or 0 => -1 (sign > 0)  */
510                                 avl_adjust_balance_factor(node, -sign);
511
512                                 /* Height is unchanged; nothing more to do.  */
513                                 return NULL;
514                         } else {
515                                 avl_adjust_balance_factor(parent, -sign);
516                                 avl_adjust_balance_factor(node, -sign);
517                         }
518                 } else {
519                         node = avl_do_double_rotate(root_ptr, node,
520                                                     parent, -sign);
521                 }
522         }
523         parent = avl_get_parent(node);
524         if (parent)
525                 *left_deleted_ret = (node == parent->left);
526         return parent;
527 }
528
529 /* Swaps node X, which must have 2 children, with its in-order successor, then
530  * unlinks node X.  Returns the parent of X just before unlinking, without its
531  * balance factor having been updated to account for the unlink.  */
532 static AVL_INLINE struct avl_tree_node *
533 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
534                              struct avl_tree_node *X,
535                              bool *left_deleted_ret)
536 {
537         struct avl_tree_node *Y, *P, *ret;
538
539         Y = X->right;
540         if (!Y->left) {
541                 /*
542                  *     P?           P?           P?
543                  *     |            |            |
544                  *     X            Y            Y
545                  *    / \          / \          / \
546                  *   A   Y    =>  A   X    =>  A   B?
547                  *      / \          / \
548                  *    (0)  B?      (0)  B?
549                  *
550                  * [ X removed, Y returned ]
551                  */
552                 ret = Y;
553                 *left_deleted_ret = false;
554         } else {
555                 struct avl_tree_node *Q;
556
557                 do {
558                         Q = Y;
559                         Y = Y->left;
560                 } while (Y->left);
561
562                 /*
563                  *     P?           P?           P?
564                  *     |            |            |
565                  *     X            Y            Y
566                  *    / \          / \          / \
567                  *   A   ...  =>  A  ...   =>  A  ...
568                  *       |            |            |
569                  *       Q            Q            Q
570                  *      /            /            /
571                  *     Y            X            B?
572                  *    / \          / \
573                  *  (0)  B?      (0)  B?
574                  *
575                  *
576                  * [ X removed, Q returned ]
577                  */
578
579                 Q->left = Y->right;
580                 if (Q->left)
581                         avl_set_parent(Q->left, Q);
582                 Y->right = X->right;
583                 avl_set_parent(X->right, Y);
584                 ret = Q;
585                 *left_deleted_ret = true;
586         }
587
588         Y->left = X->left;
589         avl_set_parent(X->left, Y);
590
591         Y->parent_balance = X->parent_balance;
592         P = avl_get_parent(X);
593         if (P) {
594                 if (P->left == X)
595                         P->left = Y;
596                 else
597                         P->right = Y;
598         } else {
599                 *root_ptr = Y;
600         }
601
602         return ret;
603 }
604
605 /* Removes the specified @node from the AVL tree.  @root_ptr must point to the
606  * pointer to the root node of the tree; *root_ptr may change if the tree is
607  * rebalanced.
608  *
609  * This *only* removes the node and rebalances the tree; it does not free
610  * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
611 void
612 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
613 {
614         struct avl_tree_node *child, *parent;
615         bool left_deleted;
616
617         if (node->left && node->right) {
618                 parent = avl_tree_swap_with_successor(root_ptr, node,
619                                                       &left_deleted);
620         } else {
621                 /* Unlink @node  */
622                 child = node->left ? node->left : node->right;
623                 parent = avl_get_parent(node);
624                 if (parent) {
625                         if (node == parent->left) {
626                                 parent->left = child;
627                                 left_deleted = true;
628                         } else {
629                                 parent->right = child;
630                                 left_deleted = false;
631                         }
632                 } else {
633                         *root_ptr = child;
634                 }
635                 if (child)
636                         avl_set_parent(child, parent);
637                 if (!parent)
638                         return;
639         }
640
641         /* Rebalance the tree  */
642         do {
643                 if (left_deleted)
644                         parent = avl_handle_subtree_shrink(root_ptr, parent,
645                                                            +1, &left_deleted);
646                 else
647                         parent = avl_handle_subtree_shrink(root_ptr, parent,
648                                                            -1, &left_deleted);
649         } while (parent);
650 }