2 * avl_tree.c - intrusive, nonrecursive AVL tree data structure (self-balancing
3 * binary search tree), implementation file
5 * The following copying information applies to this specific source code file:
7 * Written in 2014-2016 by Eric Biggers <ebiggers3@gmail.com>
9 * To the extent possible under law, the author(s) have dedicated all copyright
10 * and related and neighboring rights to this software to the public domain
11 * worldwide via the Creative Commons Zero 1.0 Universal Public Domain
12 * Dedication (the "CC0").
14 * This software is distributed in the hope that it will be useful, but WITHOUT
15 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
16 * FOR A PARTICULAR PURPOSE. See the CC0 for more details.
18 * You should have received a copy of the CC0 along with this software; if not
19 * see <http://creativecommons.org/publicdomain/zero/1.0/>.
26 #include "wimlib/avl_tree.h"
28 /* Returns the left child (sign < 0) or the right child (sign > 0) of the
29 * specified AVL tree node.
30 * Note: for all calls of this, 'sign' is constant at compilation time,
31 * so the compiler can remove the conditional. */
32 static AVL_INLINE struct avl_tree_node *
33 avl_get_child(const struct avl_tree_node *parent, int sign)
41 static AVL_INLINE struct avl_tree_node *
42 avl_tree_first_or_last_in_order(const struct avl_tree_node *root, int sign)
44 const struct avl_tree_node *first = root;
47 while (avl_get_child(first, +sign))
48 first = avl_get_child(first, +sign);
49 return (struct avl_tree_node *)first;
52 /* Starts an in-order traversal of the tree: returns the least-valued node, or
53 * NULL if the tree is empty. */
54 struct avl_tree_node *
55 avl_tree_first_in_order(const struct avl_tree_node *root)
57 return avl_tree_first_or_last_in_order(root, -1);
60 /* Starts a *reverse* in-order traversal of the tree: returns the
61 * greatest-valued node, or NULL if the tree is empty. */
62 struct avl_tree_node *
63 avl_tree_last_in_order(const struct avl_tree_node *root)
65 return avl_tree_first_or_last_in_order(root, 1);
68 static AVL_INLINE struct avl_tree_node *
69 avl_tree_next_or_prev_in_order(const struct avl_tree_node *node, int sign)
71 const struct avl_tree_node *next;
73 if (avl_get_child(node, +sign))
74 for (next = avl_get_child(node, +sign);
75 avl_get_child(next, -sign);
76 next = avl_get_child(next, -sign))
79 for (next = avl_get_parent(node);
80 next && node == avl_get_child(next, +sign);
81 node = next, next = avl_get_parent(next))
83 return (struct avl_tree_node *)next;
86 /* Continues an in-order traversal of the tree: returns the next-greatest-valued
87 * node, or NULL if there is none. */
88 struct avl_tree_node *
89 avl_tree_next_in_order(const struct avl_tree_node *node)
91 return avl_tree_next_or_prev_in_order(node, 1);
94 /* Continues a *reverse* in-order traversal of the tree: returns the
95 * previous-greatest-valued node, or NULL if there is none. */
96 struct avl_tree_node *
97 avl_tree_prev_in_order(const struct avl_tree_node *node)
99 return avl_tree_next_or_prev_in_order(node, -1);
102 /* Starts a postorder traversal of the tree. */
103 struct avl_tree_node *
104 avl_tree_first_in_postorder(const struct avl_tree_node *root)
106 const struct avl_tree_node *first = root;
109 while (first->left || first->right)
110 first = first->left ? first->left : first->right;
112 return (struct avl_tree_node *)first;
115 /* Continues a postorder traversal of the tree. @prev will not be deferenced as
116 * it's allowed that its memory has been freed; @prev_parent must be its saved
117 * parent node. Returns NULL if there are no more nodes (i.e. @prev was the
118 * root of the tree). */
119 struct avl_tree_node *
120 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
121 const struct avl_tree_node *prev_parent)
123 const struct avl_tree_node *next = prev_parent;
125 if (next && prev == next->left && next->right)
126 for (next = next->right;
127 next->left || next->right;
128 next = next->left ? next->left : next->right)
130 return (struct avl_tree_node *)next;
133 /* Sets the left child (sign < 0) or the right child (sign > 0) of the
134 * specified AVL tree node.
135 * Note: for all calls of this, 'sign' is constant at compilation time,
136 * so the compiler can remove the conditional. */
137 static AVL_INLINE void
138 avl_set_child(struct avl_tree_node *parent, int sign,
139 struct avl_tree_node *child)
142 parent->left = child;
144 parent->right = child;
147 /* Sets the parent and balance factor of the specified AVL tree node. */
148 static AVL_INLINE void
149 avl_set_parent_balance(struct avl_tree_node *node, struct avl_tree_node *parent,
152 node->parent_balance = (uintptr_t)parent | (balance_factor + 1);
155 /* Sets the parent of the specified AVL tree node. */
156 static AVL_INLINE void
157 avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
159 node->parent_balance = (uintptr_t)parent | (node->parent_balance & 3);
162 /* Returns the balance factor of the specified AVL tree node --- that is, the
163 * height of its right subtree minus the height of its left subtree. */
164 static AVL_INLINE int
165 avl_get_balance_factor(const struct avl_tree_node *node)
167 return (int)(node->parent_balance & 3) - 1;
170 /* Adds @amount to the balance factor of the specified AVL tree node.
171 * The caller must ensure this still results in a valid balance factor
173 static AVL_INLINE void
174 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
176 node->parent_balance += amount;
179 static AVL_INLINE void
180 avl_replace_child(struct avl_tree_node **root_ptr,
181 struct avl_tree_node *parent,
182 struct avl_tree_node *old_child,
183 struct avl_tree_node *new_child)
186 if (old_child == parent->left)
187 parent->left = new_child;
189 parent->right = new_child;
191 *root_ptr = new_child;
196 * Template for performing a single rotation ---
198 * sign > 0: Rotate clockwise (right) rooted at A:
208 * (nodes marked with ? may not exist)
210 * sign < 0: Rotate counterclockwise (left) rooted at A:
220 * This updates pointers but not balance factors!
222 static AVL_INLINE void
223 avl_rotate(struct avl_tree_node ** const root_ptr,
224 struct avl_tree_node * const A, const int sign)
226 struct avl_tree_node * const B = avl_get_child(A, -sign);
227 struct avl_tree_node * const E = avl_get_child(B, +sign);
228 struct avl_tree_node * const P = avl_get_parent(A);
230 avl_set_child(A, -sign, E);
231 avl_set_parent(A, B);
233 avl_set_child(B, +sign, A);
234 avl_set_parent(B, P);
237 avl_set_parent(E, A);
239 avl_replace_child(root_ptr, P, A, B);
243 * Template for performing a double rotation ---
245 * sign > 0: Rotate counterclockwise (left) rooted at B, then
246 * clockwise (right) rooted at A:
252 * B C? => E C? => B A
254 * D? E B G? D? F?G? C?
258 * (nodes marked with ? may not exist)
260 * sign < 0: Rotate clockwise (right) rooted at B, then
261 * counterclockwise (left) rooted at A:
267 * C? B => C? E => A B
269 * E D? G? B C? G?F? D?
273 * Returns a pointer to E and updates balance factors. Except for those
274 * two things, this function is equivalent to:
275 * avl_rotate(root_ptr, B, -sign);
276 * avl_rotate(root_ptr, A, +sign);
278 * See comment in avl_handle_subtree_growth() for explanation of balance
281 static AVL_INLINE struct avl_tree_node *
282 avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
283 struct avl_tree_node * const B,
284 struct avl_tree_node * const A, const int sign)
286 struct avl_tree_node * const E = avl_get_child(B, +sign);
287 struct avl_tree_node * const F = avl_get_child(E, -sign);
288 struct avl_tree_node * const G = avl_get_child(E, +sign);
289 struct avl_tree_node * const P = avl_get_parent(A);
290 const int e = avl_get_balance_factor(E);
292 avl_set_child(A, -sign, G);
293 avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
295 avl_set_child(B, +sign, F);
296 avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
298 avl_set_child(E, +sign, A);
299 avl_set_child(E, -sign, B);
300 avl_set_parent_balance(E, P, 0);
303 avl_set_parent(G, A);
306 avl_set_parent(F, B);
308 avl_replace_child(root_ptr, P, A, E);
314 * This function handles the growth of a subtree due to an insertion.
317 * Location of the tree's root pointer.
320 * A subtree that has increased in height by 1 due to an insertion.
323 * Parent of @node; must not be NULL.
326 * -1 if @node is the left child of @parent;
327 * +1 if @node is the right child of @parent.
329 * This function will adjust @parent's balance factor, then do a (single
330 * or double) rotation if necessary. The return value will be %true if
331 * the full AVL tree is now adequately balanced, or %false if the subtree
332 * rooted at @parent is now adequately balanced but has increased in
333 * height by 1, so the caller should continue up the tree.
335 * Note that if %false is returned, no rotation will have been done.
336 * Indeed, a single node insertion cannot require that more than one
337 * (single or double) rotation be done.
339 static AVL_INLINE bool
340 avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
341 struct avl_tree_node * const node,
342 struct avl_tree_node * const parent,
345 int old_balance_factor, new_balance_factor;
347 old_balance_factor = avl_get_balance_factor(parent);
349 if (old_balance_factor == 0) {
350 avl_adjust_balance_factor(parent, sign);
351 /* @parent is still sufficiently balanced (-1 or +1
352 * balance factor), but must have increased in height.
353 * Continue up the tree. */
357 new_balance_factor = old_balance_factor + sign;
359 if (new_balance_factor == 0) {
360 avl_adjust_balance_factor(parent, sign);
361 /* @parent is now perfectly balanced (0 balance factor).
362 * It cannot have increased in height, so there is
363 * nothing more to do. */
367 /* @parent is too left-heavy (new_balance_factor == -2) or
368 * too right-heavy (new_balance_factor == +2). */
370 /* Test whether @node is left-heavy (-1 balance factor) or
371 * right-heavy (+1 balance factor).
372 * Note that it cannot be perfectly balanced (0 balance factor)
373 * because here we are under the invariant that @node has
374 * increased in height due to the insertion. */
375 if (sign * avl_get_balance_factor(node) > 0) {
377 /* @node (B below) is heavy in the same direction @parent
378 * (A below) is heavy.
380 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
381 * The comment, diagram, and equations below assume sign < 0.
382 * The other case is symmetric!
383 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
385 * Do a clockwise rotation rooted at @parent (A below):
395 * Before the rotation:
398 * Let x = height(C). Then:
402 * max(height(F), height(G)) = x.
404 * After the rotation:
405 * height(D) = max(height(F), height(G)) + 1
407 * height(A) = max(height(E), height(C)) + 1
408 * = max(x, x) + 1 = x + 1
412 avl_rotate(root_ptr, parent, -sign);
414 /* Equivalent to setting @parent's balance factor to 0. */
415 avl_adjust_balance_factor(parent, -sign); /* A */
417 /* Equivalent to setting @node's balance factor to 0. */
418 avl_adjust_balance_factor(node, -sign); /* B */
420 /* @node (B below) is heavy in the direction opposite
421 * from the direction @parent (A below) is heavy.
423 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
424 * The comment, diagram, and equations below assume sign < 0.
425 * The other case is symmetric!
426 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
428 * Do a counterblockwise rotation rooted at @node (B below),
429 * then a clockwise rotation rooted at @parent (A below):
433 * B C? => E C? => B A
435 * D? E B G? D? F?G? C?
439 * Before the rotation:
442 * Let x = height(C). Then:
446 * max(height(F), height(G)) = x
448 * After both rotations:
449 * height(A) = max(height(G), height(C)) + 1
451 * balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
452 * height(B) = max(height(D), height(F)) + 1
454 * balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
459 avl_do_double_rotate(root_ptr, node, parent, -sign);
462 /* Height after rotation is unchanged; nothing more to do. */
466 /* Rebalance the tree after insertion of the specified node. */
468 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
469 struct avl_tree_node *inserted)
471 struct avl_tree_node *node, *parent;
474 inserted->left = NULL;
475 inserted->right = NULL;
479 /* Adjust balance factor of new node's parent.
480 * No rotation will need to be done at this level. */
482 parent = avl_get_parent(node);
486 if (node == parent->left)
487 avl_adjust_balance_factor(parent, -1);
489 avl_adjust_balance_factor(parent, +1);
491 if (avl_get_balance_factor(parent) == 0)
492 /* @parent did not change in height. Nothing more to do. */
495 /* The subtree rooted at @parent increased in height by 1. */
498 /* Adjust balance factor of next ancestor. */
501 parent = avl_get_parent(node);
505 /* The subtree rooted at @node has increased in height by 1. */
506 if (node == parent->left)
507 done = avl_handle_subtree_growth(root_ptr, node,
510 done = avl_handle_subtree_growth(root_ptr, node,
516 * This function handles the shrinkage of a subtree due to a deletion.
519 * Location of the tree's root pointer.
522 * A node in the tree, exactly one of whose subtrees has decreased
523 * in height by 1 due to a deletion. (This includes the case where
524 * one of the child pointers has become NULL, since we can consider
525 * the "NULL" subtree to have a height of 0.)
528 * +1 if the left subtree of @parent has decreased in height by 1;
529 * -1 if the right subtree of @parent has decreased in height by 1.
532 * If the return value is not NULL, this will be set to %true if the
533 * left subtree of the returned node has decreased in height by 1,
534 * or %false if the right subtree of the returned node has decreased
537 * This function will adjust @parent's balance factor, then do a (single
538 * or double) rotation if necessary. The return value will be NULL if
539 * the full AVL tree is now adequately balanced, or a pointer to the
540 * parent of @parent if @parent is now adequately balanced but has
541 * decreased in height by 1. Also in the latter case, *left_deleted_ret
544 static AVL_INLINE struct avl_tree_node *
545 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
546 struct avl_tree_node *parent,
548 bool * const left_deleted_ret)
550 struct avl_tree_node *node;
551 int old_balance_factor, new_balance_factor;
553 old_balance_factor = avl_get_balance_factor(parent);
555 if (old_balance_factor == 0) {
556 /* Prior to the deletion, the subtree rooted at
557 * @parent was perfectly balanced. It's now
558 * unbalanced by 1, but that's okay and its height
559 * hasn't changed. Nothing more to do. */
560 avl_adjust_balance_factor(parent, sign);
564 new_balance_factor = old_balance_factor + sign;
566 if (new_balance_factor == 0) {
567 /* The subtree rooted at @parent is now perfectly
568 * balanced, whereas before the deletion it was
569 * unbalanced by 1. Its height must have decreased
570 * by 1. No rotation is needed at this location,
571 * but continue up the tree. */
572 avl_adjust_balance_factor(parent, sign);
575 /* @parent is too left-heavy (new_balance_factor == -2) or
576 * too right-heavy (new_balance_factor == +2). */
578 node = avl_get_child(parent, sign);
580 /* The rotations below are similar to those done during
581 * insertion (see avl_handle_subtree_growth()), so full
582 * comments are not provided. The only new case is the
583 * one where @node has a balance factor of 0, and that is
586 if (sign * avl_get_balance_factor(node) >= 0) {
588 avl_rotate(root_ptr, parent, -sign);
590 if (avl_get_balance_factor(node) == 0) {
592 * @node (B below) is perfectly balanced.
594 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
595 * The comment, diagram, and equations
596 * below assume sign < 0. The other case
598 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
600 * Do a clockwise rotation rooted at
611 * Before the rotation:
614 * Let x = height(C). Then:
618 * max(height(F), height(G)) = x.
620 * After the rotation:
621 * height(D) = max(height(F), height(G)) + 1
623 * height(A) = max(height(E), height(C)) + 1
624 * = max(x + 1, x) + 1 = x + 2
629 /* A: -2 => -1 (sign < 0)
630 * or +2 => +1 (sign > 0)
631 * No change needed --- that's the same as
632 * old_balance_factor. */
634 /* B: 0 => +1 (sign < 0)
635 * or 0 => -1 (sign > 0) */
636 avl_adjust_balance_factor(node, -sign);
638 /* Height is unchanged; nothing more to do. */
641 avl_adjust_balance_factor(parent, -sign);
642 avl_adjust_balance_factor(node, -sign);
645 node = avl_do_double_rotate(root_ptr, node,
649 parent = avl_get_parent(node);
651 *left_deleted_ret = (node == parent->left);
655 /* Swaps node X, which must have 2 children, with its in-order successor, then
656 * unlinks node X. Returns the parent of X just before unlinking, without its
657 * balance factor having been updated to account for the unlink. */
658 static AVL_INLINE struct avl_tree_node *
659 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
660 struct avl_tree_node *X,
661 bool *left_deleted_ret)
663 struct avl_tree_node *Y, *ret;
676 * [ X unlinked, Y returned ]
679 *left_deleted_ret = false;
681 struct avl_tree_node *Q;
693 * A ... => A ... => A ...
702 * [ X unlinked, Q returned ]
707 avl_set_parent(Q->left, Q);
709 avl_set_parent(X->right, Y);
711 *left_deleted_ret = true;
715 avl_set_parent(X->left, Y);
717 Y->parent_balance = X->parent_balance;
718 avl_replace_child(root_ptr, avl_get_parent(X), X, Y);
724 * Removes an item from the specified AVL tree.
727 * Location of the AVL tree's root pointer. Indirection is needed
728 * because the root node may change if the tree needed to be rebalanced
729 * because of the deletion or if @node was the root node.
732 * Pointer to the `struct avl_tree_node' embedded in the item to
733 * remove from the tree.
735 * Note: This function *only* removes the node and rebalances the tree.
736 * It does not free any memory.
739 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
741 struct avl_tree_node *parent;
742 bool left_deleted = false;
744 if (node->left && node->right) {
745 /* @node is fully internal, with two children. Swap it
746 * with its in-order successor (which must exist in the
747 * right subtree of @node and can have, at most, a right
748 * child), then unlink @node. */
749 parent = avl_tree_swap_with_successor(root_ptr, node,
751 /* @parent is now the parent of what was @node's in-order
752 * successor. It cannot be NULL, since @node itself was
753 * an ancestor of its in-order successor.
754 * @left_deleted has been set to %true if @node's
755 * in-order successor was the left child of @parent,
756 * otherwise %false. */
758 struct avl_tree_node *child;
760 /* @node is missing at least one child. Unlink it. Set
761 * @parent to @node's parent, and set @left_deleted to
762 * reflect which child of @parent @node was. Or, if
763 * @node was the root node, simply update the root node
765 child = node->left ? node->left : node->right;
766 parent = avl_get_parent(node);
768 if (node == parent->left) {
769 parent->left = child;
772 parent->right = child;
773 left_deleted = false;
776 avl_set_parent(child, parent);
779 avl_set_parent(child, parent);
785 /* Rebalance the tree. */
788 parent = avl_handle_subtree_shrink(root_ptr, parent,
791 parent = avl_handle_subtree_shrink(root_ptr, parent,