95e0e187923590129bbee8375e4881c13c70d501
[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         /* Height of @node subtree of @parent has increased by 1.
278          * Adjust @parent's balance factor and check whether rotations need to
279          * be done.  */
280
281         int old_balance_factor, new_balance_factor;
282
283         old_balance_factor = avl_get_balance_factor(parent);
284
285         if (old_balance_factor == 0) {
286                 /* Parent has increased in height but is still sufficiently
287                  * balanced.  Continue up the tree.  */
288                 avl_adjust_balance_factor(parent, sign);
289                 return false;
290         }
291
292         new_balance_factor = old_balance_factor + sign;
293
294         if (new_balance_factor == 0) {
295                 /* Parent is balanced; nothing more to to.  */
296                 avl_adjust_balance_factor(parent, sign);
297                 return true;
298         }
299
300         /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
301          * The other case is symmetric --- that is, the rotations done are the
302          * the mirror images, all the balance factors are inverted, and left and
303          * right pointers are otherwise reversed.  */
304
305         /* Parent is left-heavy (balance_factor == -2).  */
306
307         if (sign * avl_get_balance_factor(node) > 0) {
308
309                 /* Child node (@node, also B below) is also left-heavy.
310                  * It must have balance_factor == -1.
311                  * Do a clockwise ("right") rotation rooted at
312                  * @parent (A below):
313                  *
314                  *           A              B
315                  *          / \           /   \
316                  *         B   C   =>    D     A
317                  *        / \           / \   / \
318                  *       D   E         F   G E   C
319                  *      / \
320                  *     F   G
321                  *
322                  * Before the rotation:
323                  *      balance(A) = -2
324                  *      balance(B) = -1
325                  * Let x = height(C).  Then:
326                  *      height(B) = x + 2
327                  *      height(D) = x + 1
328                  *      height(E) = x
329                  *      max(height(F), height(G)) = x.
330                  *
331                  * After the rotation:
332                  *      height(D) = max(height(F), height(G)) + 1
333                  *                = x + 1
334                  *      height(A) = max(height(E), height(C)) + 1
335                  *                = max(x, x) + 1 = x + 1
336                  *      balance(B) = 0
337                  *      balance(A) = 0
338                  */
339
340                 avl_rotate(root_ptr, parent, -sign);
341                 avl_adjust_balance_factor(parent, -sign); /* A */
342                 avl_adjust_balance_factor(node, -sign);   /* B */
343         } else {
344                 /* Child node (@node, also B below) is right-heavy.
345                  * It must have balance_factor == +1.
346                  * Do a counterclockwise ("left") rotation rooted at child node
347                  * (B below), then a clockwise ("right") rotation rooted at
348                  * parent node (A below).
349                  *
350                  *           A             A           E
351                  *          / \           / \        /   \
352                  *         B   C   =>    E   C  =>  B     A
353                  *        / \           / \        / \   / \
354                  *       D   E         B   G      D   F G   C
355                  *          / \       / \
356                  *         F   G     D   F
357                  *
358                  * Before the rotation:
359                  *      balance(A) = -2
360                  *      balance(B) = +1
361                  * Let x = height(C).  Then:
362                  *      height(B) = x + 2
363                  *      height(E) = x + 1
364                  *      height(D) = x
365                  *      max(height(F), height(G)) = x
366                  *
367                  * After both rotations:
368                  *      height(A) = max(height(G), height(C)) + 1
369                  *                = x + 1
370                  *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
371                  *      height(B) = max(height(D), height(F)) + 1
372                  *                = x + 1
373                  *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
374                  *
375                  *      height(E) = x + 2
376                  *      balance(E) = 0
377                  */
378                 avl_do_double_rotate(root_ptr, node, parent, -sign);
379         }
380         return true;
381 }
382
383 /* Rebalance the tree after insertion of the specified node.  */
384 void
385 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
386                                 struct avl_tree_node *inserted)
387 {
388         struct avl_tree_node *node, *parent;
389         bool done = false;
390
391         inserted->left = NULL;
392         inserted->right = NULL;
393
394         for (node = inserted, parent = avl_get_parent(node);
395              parent && !done;
396              node = parent, parent = avl_get_parent(parent))
397         {
398                 /* Height of @node subtree has increased by 1  */
399
400                 if (node == parent->left)
401                         done = avl_handle_subtree_growth(root_ptr, node,
402                                                          parent, -1);
403                 else
404                         done = avl_handle_subtree_growth(root_ptr, node,
405                                                          parent, +1);
406         }
407 }
408
409 static AVL_INLINE struct avl_tree_node *
410 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
411                           struct avl_tree_node *parent,
412                           const int sign,
413                           bool * const left_deleted_ret)
414 {
415         struct avl_tree_node *node;
416
417         const int old_balance_factor = avl_get_balance_factor(parent);
418         const int new_balance_factor = old_balance_factor + sign;
419
420         if (old_balance_factor == 0) {
421                 /* Prior to the deletion, the subtree rooted at
422                  * @parent was perfectly balanced.  It's now
423                  * unbalanced by 1, but that's okay and its height
424                  * hasn't changed.  Nothing more to do.  */
425                 avl_adjust_balance_factor(parent, sign);
426                 return NULL;
427         } else if (new_balance_factor == 0) {
428                 /* The subtree rooted at @parent is now perfectly
429                  * balanced, whereas before the deletion it was
430                  * unbalanced by 1.  Its height must have decreased
431                  * by 1.  No rotation is needed at this location,
432                  * but continue up the tree.  */
433                 avl_adjust_balance_factor(parent, sign);
434                 node = parent;
435         } else {
436                 /* The subtree rooted at @parent is now significantly
437                  * unbalanced (by 2 in some direction).  */
438                 node = avl_get_child(parent, sign);
439
440                 /* The rotations below are similar to those done during
441                  * insertion.  The only new case is the one where the
442                  * child node has a balance factor of 0.  */
443
444                 if (sign * avl_get_balance_factor(node) >= 0) {
445                         avl_rotate(root_ptr, parent, -sign);
446
447                         if (avl_get_balance_factor(node) == 0) {
448                                 /* Child node (@node, also B below) is balanced.
449                                  * Do a clockwise ("right") rotation rooted at
450                                  * @parent (A below):
451                                  *
452                                  *           A              B
453                                  *          / \           /   \
454                                  *         B   C   =>    D     A
455                                  *        / \           / \   / \
456                                  *       D   E         F   G E   C
457                                  *      / \
458                                  *     F   G
459                                  *
460                                  * Before the rotation:
461                                  *      balance(A) = -2
462                                  *      balance(B) =  0
463                                  * Let x = height(C).  Then:
464                                  *      height(B) = x + 2
465                                  *      height(D) = x + 1
466                                  *      height(E) = x + 1
467                                  *      max(height(F), height(G)) = x.
468                                  *
469                                  * After the rotation:
470                                  *      height(D) = max(height(F), height(G)) + 1
471                                  *                = x + 1
472                                  *      height(A) = max(height(E), height(C)) + 1
473                                  *                = max(x + 1, x) + 1 = x + 2
474                                  *      balance(A) = -1
475                                  *      balance(B) = +1
476                                  */
477
478                                 /* A: -2 => -1 (sign < 0)
479                                  * or +2 => +1 (sign > 0)
480                                  * No change needed --- that's the same as
481                                  * old_balance_factor.  */
482
483                                 /* B: 0 => +1 (sign < 0)
484                                  * or 0 => -1 (sign > 0)  */
485                                 avl_adjust_balance_factor(node, -sign);
486
487                                 /* Height is unchanged; nothing more to do.  */
488                                 return NULL;
489                         } else {
490                                 avl_adjust_balance_factor(parent, -sign);
491                                 avl_adjust_balance_factor(node, -sign);
492                         }
493                 } else {
494                         node = avl_do_double_rotate(root_ptr, node,
495                                                     parent, -sign);
496                 }
497         }
498         parent = avl_get_parent(node);
499         if (parent)
500                 *left_deleted_ret = (node == parent->left);
501         return parent;
502 }
503
504 /* Swaps node X, which must have 2 children, with its in-order successor, then
505  * unlinks node X.  Returns the parent of X just before unlinking, without its
506  * balance factor having been updated to account for the unlink.  */
507 static AVL_INLINE struct avl_tree_node *
508 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
509                              struct avl_tree_node *X,
510                              bool *left_deleted_ret)
511 {
512         struct avl_tree_node *Y, *P, *ret;
513
514         Y = X->right;
515         if (!Y->left) {
516                 /*
517                  *     P?           P?           P?
518                  *     |            |            |
519                  *     X            Y            Y
520                  *    / \          / \          / \
521                  *   A   Y    =>  A   X    =>  A   B?
522                  *      / \          / \
523                  *    (0)  B?      (0)  B?
524                  *
525                  * [ X removed, Y returned ]
526                  */
527                 ret = Y;
528                 *left_deleted_ret = false;
529         } else {
530                 struct avl_tree_node *Q;
531
532                 do {
533                         Q = Y;
534                         Y = Y->left;
535                 } while (Y->left);
536
537                 /*
538                  *     P?           P?           P?
539                  *     |            |            |
540                  *     X            Y            Y
541                  *    / \          / \          / \
542                  *   A   ...  =>  A  ...   =>  A  ...
543                  *       |            |            |
544                  *       Q            Q            Q
545                  *      /            /            /
546                  *     Y            X            B?
547                  *    / \          / \
548                  *  (0)  B?      (0)  B?
549                  *
550                  *
551                  * [ X removed, Q returned ]
552                  */
553
554                 Q->left = Y->right;
555                 if (Q->left)
556                         avl_set_parent(Q->left, Q);
557                 Y->right = X->right;
558                 avl_set_parent(X->right, Y);
559                 ret = Q;
560                 *left_deleted_ret = true;
561         }
562
563         Y->left = X->left;
564         avl_set_parent(X->left, Y);
565
566         Y->parent_balance = X->parent_balance;
567         P = avl_get_parent(X);
568         if (P) {
569                 if (P->left == X)
570                         P->left = Y;
571                 else
572                         P->right = Y;
573         } else {
574                 *root_ptr = Y;
575         }
576
577         return ret;
578 }
579
580 /* Removes the specified @node from the AVL tree.  @root_ptr must point to the
581  * pointer to the root node of the tree; *root_ptr may change if the tree is
582  * rebalanced.
583  *
584  * This *only* removes the node and rebalances the tree; it does not free
585  * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
586 void
587 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
588 {
589         struct avl_tree_node *child, *parent;
590         bool left_deleted;
591
592         if (node->left && node->right) {
593                 parent = avl_tree_swap_with_successor(root_ptr, node,
594                                                       &left_deleted);
595         } else {
596                 /* Unlink @node  */
597                 child = node->left ? node->left : node->right;
598                 parent = avl_get_parent(node);
599                 if (parent) {
600                         if (node == parent->left) {
601                                 parent->left = child;
602                                 left_deleted = true;
603                         } else {
604                                 parent->right = child;
605                                 left_deleted = false;
606                         }
607                 } else {
608                         *root_ptr = child;
609                 }
610                 if (child)
611                         avl_set_parent(child, parent);
612                 if (!parent)
613                         return;
614         }
615
616         /* Rebalance the tree  */
617         do {
618                 if (left_deleted)
619                         parent = avl_handle_subtree_shrink(root_ptr, parent,
620                                                            +1, &left_deleted);
621                 else
622                         parent = avl_handle_subtree_shrink(root_ptr, parent,
623                                                            -1, &left_deleted);
624         } while (parent);
625 }