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 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 /* Starts an in-order traversal of the tree: returns the least-valued node, or
29 * NULL if the tree is empty. */
30 struct avl_tree_node *
31 avl_tree_first_in_order(const struct avl_tree_node *root)
33 const struct avl_tree_node *first = root;
38 return (struct avl_tree_node *)first;
41 /* Continues an in-order traversal of the tree: returns the next-greatest-valued
42 * node, or NULL if there is none. */
43 struct avl_tree_node *
44 avl_tree_next_in_order(const struct avl_tree_node *node)
46 const struct avl_tree_node *next;
49 for (next = node->right;
54 for (next = avl_get_parent(node);
55 next && node == next->right;
56 node = next, next = avl_get_parent(next))
58 return (struct avl_tree_node *)next;
61 /* Continues a *reverse* in-order traversal of the tree: returns the
62 * previous-greatest-valued node, or NULL if there is none. */
63 struct avl_tree_node *
64 avl_tree_prev_in_order(const struct avl_tree_node *node)
66 const struct avl_tree_node *prev;
69 for (prev = node->left;
74 for (prev = avl_get_parent(node);
75 prev && node == prev->left;
76 node = prev, prev = avl_get_parent(prev))
78 return (struct avl_tree_node *)prev;
81 /* Starts a postorder traversal of the tree. */
82 struct avl_tree_node *
83 avl_tree_first_in_postorder(const struct avl_tree_node *root)
85 const struct avl_tree_node *first = root;
88 while (first->left || first->right)
89 first = first->left ? first->left : first->right;
91 return (struct avl_tree_node *)first;
94 /* Continues a postorder traversal of the tree. @prev will not be deferenced as
95 * it's allowed that its memory has been freed; @prev_parent must be its saved
96 * parent node. Returns NULL if there are no more nodes (i.e. @prev was the
97 * root of the tree). */
98 struct avl_tree_node *
99 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
100 const struct avl_tree_node *prev_parent)
102 const struct avl_tree_node *next = prev_parent;
104 if (next && prev == next->left && next->right)
105 for (next = next->right;
106 next->left || next->right;
107 next = next->left ? next->left : next->right)
109 return (struct avl_tree_node *)next;
112 /* Returns the left child (sign < 0) or the right child (sign > 0) of the
113 * specified AVL tree node.
114 * Note: for all calls of this, 'sign' is constant at compilation time,
115 * so the compiler can remove the conditional. */
116 static forceinline struct avl_tree_node *
117 avl_get_child(const struct avl_tree_node *parent, int sign)
122 return parent->right;
125 /* Sets the left child (sign < 0) or the right child (sign > 0) of the
126 * specified AVL tree node.
127 * Note: for all calls of this, 'sign' is constant at compilation time,
128 * so the compiler can remove the conditional. */
129 static forceinline void
130 avl_set_child(struct avl_tree_node *parent, int sign,
131 struct avl_tree_node *child)
134 parent->left = child;
136 parent->right = child;
139 /* Sets the parent and balance factor of the specified AVL tree node. */
140 static forceinline void
141 avl_set_parent_balance(struct avl_tree_node *node, struct avl_tree_node *parent,
144 node->parent_balance = (uintptr_t)parent | (balance_factor + 1);
147 /* Sets the parent of the specified AVL tree node. */
148 static forceinline void
149 avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
151 node->parent_balance = (uintptr_t)parent | (node->parent_balance & 3);
154 /* Returns the balance factor of the specified AVL tree node --- that is, the
155 * height of its right subtree minus the height of its left subtree. */
156 static forceinline int
157 avl_get_balance_factor(const struct avl_tree_node *node)
159 return (int)(node->parent_balance & 3) - 1;
162 /* Adds @amount to the balance factor of the specified AVL tree node.
163 * The caller must ensure this still results in a valid balance factor
165 static forceinline void
166 avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
168 node->parent_balance += amount;
171 static forceinline void
172 avl_replace_child(struct avl_tree_node **root_ptr,
173 struct avl_tree_node *parent,
174 struct avl_tree_node *old_child,
175 struct avl_tree_node *new_child)
178 if (old_child == parent->left)
179 parent->left = new_child;
181 parent->right = new_child;
183 *root_ptr = new_child;
188 * Template for performing a single rotation ---
190 * sign > 0: Rotate clockwise (right) rooted at A:
200 * (nodes marked with ? may not exist)
202 * sign < 0: Rotate counterclockwise (left) rooted at A:
212 * This updates pointers but not balance factors!
214 static forceinline void
215 avl_rotate(struct avl_tree_node ** const root_ptr,
216 struct avl_tree_node * const A, const int sign)
218 struct avl_tree_node * const B = avl_get_child(A, -sign);
219 struct avl_tree_node * const E = avl_get_child(B, +sign);
220 struct avl_tree_node * const P = avl_get_parent(A);
222 avl_set_child(A, -sign, E);
223 avl_set_parent(A, B);
225 avl_set_child(B, +sign, A);
226 avl_set_parent(B, P);
229 avl_set_parent(E, A);
231 avl_replace_child(root_ptr, P, A, B);
235 * Template for performing a double rotation ---
237 * sign > 0: Rotate counterclockwise (left) rooted at B, then
238 * clockwise (right) rooted at A:
244 * B C? => E C? => B A
246 * D? E B G? D? F?G? C?
250 * (nodes marked with ? may not exist)
252 * sign < 0: Rotate clockwise (right) rooted at B, then
253 * counterclockwise (left) rooted at A:
259 * C? B => C? E => A B
261 * E D? G? B C? G?F? D?
265 * Returns a pointer to E and updates balance factors. Except for those
266 * two things, this function is equivalent to:
267 * avl_rotate(root_ptr, B, -sign);
268 * avl_rotate(root_ptr, A, +sign);
270 * See comment in avl_handle_subtree_growth() for explanation of balance
273 static forceinline struct avl_tree_node *
274 avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
275 struct avl_tree_node * const B,
276 struct avl_tree_node * const A, const int sign)
278 struct avl_tree_node * const E = avl_get_child(B, +sign);
279 struct avl_tree_node * const F = avl_get_child(E, -sign);
280 struct avl_tree_node * const G = avl_get_child(E, +sign);
281 struct avl_tree_node * const P = avl_get_parent(A);
282 const int e = avl_get_balance_factor(E);
284 avl_set_child(A, -sign, G);
285 avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
287 avl_set_child(B, +sign, F);
288 avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
290 avl_set_child(E, +sign, A);
291 avl_set_child(E, -sign, B);
292 avl_set_parent_balance(E, P, 0);
295 avl_set_parent(G, A);
298 avl_set_parent(F, B);
300 avl_replace_child(root_ptr, P, A, E);
306 * This function handles the growth of a subtree due to an insertion.
309 * Location of the tree's root pointer.
312 * A subtree that has increased in height by 1 due to an insertion.
315 * Parent of @node; must not be NULL.
318 * -1 if @node is the left child of @parent;
319 * +1 if @node is the right child of @parent.
321 * This function will adjust @parent's balance factor, then do a (single
322 * or double) rotation if necessary. The return value will be %true if
323 * the full AVL tree is now adequately balanced, or %false if the subtree
324 * rooted at @parent is now adequately balanced but has increased in
325 * height by 1, so the caller should continue up the tree.
327 * Note that if %false is returned, no rotation will have been done.
328 * Indeed, a single node insertion cannot require that more than one
329 * (single or double) rotation be done.
331 static forceinline bool
332 avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
333 struct avl_tree_node * const node,
334 struct avl_tree_node * const parent,
337 int old_balance_factor, new_balance_factor;
339 old_balance_factor = avl_get_balance_factor(parent);
341 if (old_balance_factor == 0) {
342 avl_adjust_balance_factor(parent, sign);
343 /* @parent is still sufficiently balanced (-1 or +1
344 * balance factor), but must have increased in height.
345 * Continue up the tree. */
349 new_balance_factor = old_balance_factor + sign;
351 if (new_balance_factor == 0) {
352 avl_adjust_balance_factor(parent, sign);
353 /* @parent is now perfectly balanced (0 balance factor).
354 * It cannot have increased in height, so there is
355 * nothing more to do. */
359 /* @parent is too left-heavy (new_balance_factor == -2) or
360 * too right-heavy (new_balance_factor == +2). */
362 /* Test whether @node is left-heavy (-1 balance factor) or
363 * right-heavy (+1 balance factor).
364 * Note that it cannot be perfectly balanced (0 balance factor)
365 * because here we are under the invariant that @node has
366 * increased in height due to the insertion. */
367 if (sign * avl_get_balance_factor(node) > 0) {
369 /* @node (B below) is heavy in the same direction @parent
370 * (A below) is heavy.
372 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
373 * The comment, diagram, and equations below assume sign < 0.
374 * The other case is symmetric!
375 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
377 * Do a clockwise rotation rooted at @parent (A below):
387 * Before the rotation:
390 * Let x = height(C). Then:
394 * max(height(F), height(G)) = x.
396 * After the rotation:
397 * height(D) = max(height(F), height(G)) + 1
399 * height(A) = max(height(E), height(C)) + 1
400 * = max(x, x) + 1 = x + 1
404 avl_rotate(root_ptr, parent, -sign);
406 /* Equivalent to setting @parent's balance factor to 0. */
407 avl_adjust_balance_factor(parent, -sign); /* A */
409 /* Equivalent to setting @node's balance factor to 0. */
410 avl_adjust_balance_factor(node, -sign); /* B */
412 /* @node (B below) is heavy in the direction opposite
413 * from the direction @parent (A below) is heavy.
415 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
416 * The comment, diagram, and equations below assume sign < 0.
417 * The other case is symmetric!
418 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
420 * Do a counterblockwise rotation rooted at @node (B below),
421 * then a clockwise rotation rooted at @parent (A below):
425 * B C? => E C? => B A
427 * D? E B G? D? F?G? C?
431 * Before the rotation:
434 * Let x = height(C). Then:
438 * max(height(F), height(G)) = x
440 * After both rotations:
441 * height(A) = max(height(G), height(C)) + 1
443 * balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
444 * height(B) = max(height(D), height(F)) + 1
446 * balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
451 avl_do_double_rotate(root_ptr, node, parent, -sign);
454 /* Height after rotation is unchanged; nothing more to do. */
458 /* Rebalance the tree after insertion of the specified node. */
460 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
461 struct avl_tree_node *inserted)
463 struct avl_tree_node *node, *parent;
466 inserted->left = NULL;
467 inserted->right = NULL;
471 /* Adjust balance factor of new node's parent.
472 * No rotation will need to be done at this level. */
474 parent = avl_get_parent(node);
478 if (node == parent->left)
479 avl_adjust_balance_factor(parent, -1);
481 avl_adjust_balance_factor(parent, +1);
483 if (avl_get_balance_factor(parent) == 0)
484 /* @parent did not change in height. Nothing more to do. */
487 /* The subtree rooted at @parent increased in height by 1. */
490 /* Adjust balance factor of next ancestor. */
493 parent = avl_get_parent(node);
497 /* The subtree rooted at @node has increased in height by 1. */
498 if (node == parent->left)
499 done = avl_handle_subtree_growth(root_ptr, node,
502 done = avl_handle_subtree_growth(root_ptr, node,
508 * This function handles the shrinkage of a subtree due to a deletion.
511 * Location of the tree's root pointer.
514 * A node in the tree, exactly one of whose subtrees has decreased
515 * in height by 1 due to a deletion. (This includes the case where
516 * one of the child pointers has become NULL, since we can consider
517 * the "NULL" subtree to have a height of 0.)
520 * +1 if the left subtree of @parent has decreased in height by 1;
521 * -1 if the right subtree of @parent has decreased in height by 1.
524 * If the return value is not NULL, this will be set to %true if the
525 * left subtree of the returned node has decreased in height by 1,
526 * or %false if the right subtree of the returned node has decreased
529 * This function will adjust @parent's balance factor, then do a (single
530 * or double) rotation if necessary. The return value will be NULL if
531 * the full AVL tree is now adequately balanced, or a pointer to the
532 * parent of @parent if @parent is now adequately balanced but has
533 * decreased in height by 1. Also in the latter case, *left_deleted_ret
536 static forceinline struct avl_tree_node *
537 avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
538 struct avl_tree_node *parent,
540 bool * const left_deleted_ret)
542 struct avl_tree_node *node;
543 int old_balance_factor, new_balance_factor;
545 old_balance_factor = avl_get_balance_factor(parent);
547 if (old_balance_factor == 0) {
548 /* Prior to the deletion, the subtree rooted at
549 * @parent was perfectly balanced. It's now
550 * unbalanced by 1, but that's okay and its height
551 * hasn't changed. Nothing more to do. */
552 avl_adjust_balance_factor(parent, sign);
556 new_balance_factor = old_balance_factor + sign;
558 if (new_balance_factor == 0) {
559 /* The subtree rooted at @parent is now perfectly
560 * balanced, whereas before the deletion it was
561 * unbalanced by 1. Its height must have decreased
562 * by 1. No rotation is needed at this location,
563 * but continue up the tree. */
564 avl_adjust_balance_factor(parent, sign);
567 /* @parent is too left-heavy (new_balance_factor == -2) or
568 * too right-heavy (new_balance_factor == +2). */
570 node = avl_get_child(parent, sign);
572 /* The rotations below are similar to those done during
573 * insertion (see avl_handle_subtree_growth()), so full
574 * comments are not provided. The only new case is the
575 * one where @node has a balance factor of 0, and that is
578 if (sign * avl_get_balance_factor(node) >= 0) {
580 avl_rotate(root_ptr, parent, -sign);
582 if (avl_get_balance_factor(node) == 0) {
584 * @node (B below) is perfectly balanced.
586 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
587 * The comment, diagram, and equations
588 * below assume sign < 0. The other case
590 * @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
592 * Do a clockwise rotation rooted at
603 * Before the rotation:
606 * Let x = height(C). Then:
610 * max(height(F), height(G)) = x.
612 * After the rotation:
613 * height(D) = max(height(F), height(G)) + 1
615 * height(A) = max(height(E), height(C)) + 1
616 * = max(x + 1, x) + 1 = x + 2
621 /* A: -2 => -1 (sign < 0)
622 * or +2 => +1 (sign > 0)
623 * No change needed --- that's the same as
624 * old_balance_factor. */
626 /* B: 0 => +1 (sign < 0)
627 * or 0 => -1 (sign > 0) */
628 avl_adjust_balance_factor(node, -sign);
630 /* Height is unchanged; nothing more to do. */
633 avl_adjust_balance_factor(parent, -sign);
634 avl_adjust_balance_factor(node, -sign);
637 node = avl_do_double_rotate(root_ptr, node,
641 parent = avl_get_parent(node);
643 *left_deleted_ret = (node == parent->left);
647 /* Swaps node X, which must have 2 children, with its in-order successor, then
648 * unlinks node X. Returns the parent of X just before unlinking, without its
649 * balance factor having been updated to account for the unlink. */
650 static forceinline struct avl_tree_node *
651 avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
652 struct avl_tree_node *X,
653 bool *left_deleted_ret)
655 struct avl_tree_node *Y, *ret;
668 * [ X unlinked, Y returned ]
671 *left_deleted_ret = false;
673 struct avl_tree_node *Q;
685 * A ... => A ... => A ...
694 * [ X unlinked, Q returned ]
699 avl_set_parent(Q->left, Q);
701 avl_set_parent(X->right, Y);
703 *left_deleted_ret = true;
707 avl_set_parent(X->left, Y);
709 Y->parent_balance = X->parent_balance;
710 avl_replace_child(root_ptr, avl_get_parent(X), X, Y);
716 * Removes an item from the specified AVL tree.
719 * Location of the AVL tree's root pointer. Indirection is needed
720 * because the root node may change if the tree needed to be rebalanced
721 * because of the deletion or if @node was the root node.
724 * Pointer to the `struct avl_tree_node' embedded in the item to
725 * remove from the tree.
727 * Note: This function *only* removes the node and rebalances the tree.
728 * It does not free any memory.
731 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
733 struct avl_tree_node *parent;
734 bool left_deleted = false;
736 if (node->left && node->right) {
737 /* @node is fully internal, with two children. Swap it
738 * with its in-order successor (which must exist in the
739 * right subtree of @node and can have, at most, a right
740 * child), then unlink @node. */
741 parent = avl_tree_swap_with_successor(root_ptr, node,
743 /* @parent is now the parent of what was @node's in-order
744 * successor. It cannot be NULL, since @node itself was
745 * an ancestor of its in-order successor.
746 * @left_deleted has been set to %true if @node's
747 * in-order successor was the left child of @parent,
748 * otherwise %false. */
750 struct avl_tree_node *child;
752 /* @node is missing at least one child. Unlink it. Set
753 * @parent to @node's parent, and set @left_deleted to
754 * reflect which child of @parent @node was. Or, if
755 * @node was the root node, simply update the root node
757 child = node->left ? node->left : node->right;
758 parent = avl_get_parent(node);
760 if (node == parent->left) {
761 parent->left = child;
764 parent->right = child;
765 left_deleted = false;
768 avl_set_parent(child, parent);
771 avl_set_parent(child, parent);
777 /* Rebalance the tree. */
780 parent = avl_handle_subtree_shrink(root_ptr, parent,
783 parent = avl_handle_subtree_shrink(root_ptr, parent,