]> wimlib.net Git - wimlib/blobdiff - src/avl_tree.c
Replace red-black tree code with AVL tree code
[wimlib] / src / avl_tree.c
diff --git a/src/avl_tree.c b/src/avl_tree.c
new file mode 100644 (file)
index 0000000..ab269d0
--- /dev/null
@@ -0,0 +1,502 @@
+/*
+ * avl_tree.c
+ *
+ * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
+ * tree), implementation file.
+ *
+ * Author:  Eric Biggers
+ * Year:    2014
+ *
+ * This file is placed into the public domain.  You can do whatever you want
+ * with it.
+ */
+
+#include <wimlib/avl_tree.h>
+
+/* Starts an in-order traversal of the tree: returns the least-valued node, or
+ * NULL if the tree is empty.  */
+struct avl_tree_node *
+avl_tree_first_in_order(const struct avl_tree_node *root)
+{
+       const struct avl_tree_node *first = root;
+
+       if (first)
+               while (first->left)
+                       first = first->left;
+       return (struct avl_tree_node *)first;
+}
+
+/* Continues an in-order traversal of the tree: returns the next-greatest-valued
+ * node, or NULL if there is none.  */
+struct avl_tree_node *
+avl_tree_next_in_order(const struct avl_tree_node *prev)
+{
+       const struct avl_tree_node *next;
+
+       if (prev->right)
+               for (next = prev->right;
+                    next->left;
+                    next = next->left)
+                       ;
+       else
+               for (next = avl_get_parent(prev);
+                    next && prev == next->right;
+                    prev = next, next = avl_get_parent(next))
+                       ;
+       return (struct avl_tree_node *)next;
+}
+
+/* Starts a postorder traversal of the tree.  */
+struct avl_tree_node *
+avl_tree_first_in_postorder(const struct avl_tree_node *root)
+{
+       const struct avl_tree_node *first = root;
+
+       if (first)
+               while (first->left || first->right)
+                       first = first->left ? first->left : first->right;
+
+       return (struct avl_tree_node *)first;
+}
+
+/* Continues a postorder traversal of the tree.  @prev will not be deferenced as
+ * it's allowed that its memory has been freed; @prev_parent must be its saved
+ * parent node.  Returns NULL if there are no more nodes (i.e. @prev was the
+ * root of the tree).  */
+struct avl_tree_node *
+avl_tree_next_in_postorder(const struct avl_tree_node *prev,
+                          const struct avl_tree_node *prev_parent)
+{
+       const struct avl_tree_node *next = prev_parent;
+
+       if (next && prev == next->left && next->right)
+               for (next = next->right;
+                    next->left || next->right;
+                    next = next->left ? next->left : next->right)
+                       ;
+       return (struct avl_tree_node *)next;
+}
+
+/* Get the left child (sign < 0) or the right child (sign > 0)
+ * Note: for all calls of this, 'sign' is constant at compilation time and the
+ * compiler can remove the conditional.  */
+static AVL_INLINE struct avl_tree_node *
+avl_get_child(const struct avl_tree_node *parent, int sign)
+{
+       if (sign < 0)
+               return parent->left;
+       else
+               return parent->right;
+}
+
+/* Set the left child (sign < 0) or the right child (sign > 0)
+ * Note: for all calls of this, 'sign' is constant at compilation time and the
+ * compiler can remove the conditional.  */
+static AVL_INLINE void
+avl_set_child(struct avl_tree_node *parent, int sign,
+             struct avl_tree_node *child)
+{
+       if (sign < 0)
+               parent->left = child;
+       else
+               parent->right = child;
+}
+
+/* Sets the parent of the AVL tree @node to @parent.  */
+static AVL_INLINE void
+avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
+{
+       node->parent_balance =
+               (node->parent_balance & 3) | (uintptr_t)parent;
+}
+
+/* Returns the balance factor of the specified AVL tree node --- that is, the
+ * height of its right subtree minus the height of its left subtree.  */
+static AVL_INLINE int
+avl_get_balance_factor(const struct avl_tree_node *node)
+{
+       return (int)(node->parent_balance & 3) - 1;
+}
+
+/* Sets the balance factor of the specified AVL tree node.  The caller MUST
+ * ensure that this number is valid and maintains the AVL tree invariants.  */
+static AVL_INLINE void
+avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
+{
+       node->parent_balance =
+               (node->parent_balance & ~3) | (balance_factor + 1);
+}
+
+/* Increments the balance factor of the specified AVL tree @node by the
+ * specified @amount.  The caller MUST ensure that this number is valid and
+ * maintains the AVL tree invariants.  */
+static AVL_INLINE void
+avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
+{
+       node->parent_balance += amount;
+}
+
+/*
+ * Template for performing a rotation ---
+ *
+ * Clockwise/Right (sign > 0):
+ *
+ *           X             X
+ *           |             |
+ *           A             B
+ *          / \           / \
+ *         B   C   =>    D   A
+ *        / \               / \
+ *       D   E             E   C
+ *
+ * Counterclockwise/Left (sign < 0)- same, except left and right are reversed:
+ *
+ *           X             X
+ *           |             |
+ *           A             B
+ *          / \           / \
+ *         C   B   =>    A   D
+ *            / \       / \
+ *           E   D     C   E
+ *
+ * This updates pointers but not balance factors!
+ */
+static AVL_INLINE void
+avl_rotate(struct avl_tree_node * const A, const int sign)
+{
+       struct avl_tree_node * const B = avl_get_child(A, -sign);
+       struct avl_tree_node * const C = avl_get_child(A, +sign);
+       struct avl_tree_node * const E = avl_get_child(B, +sign);
+       struct avl_tree_node * const X = avl_get_parent(A);
+
+       avl_set_child(A, -sign, E);
+       avl_set_child(A, +sign, C);
+       avl_set_parent(A, B);
+
+       avl_set_child(B, +sign, A);
+       avl_set_parent(B, X);
+
+       if (E)
+               avl_set_parent(E, A);
+
+       if (X) {
+               if (avl_get_child(X, +sign) == A)
+                       avl_set_child(X, +sign, B);
+               else
+                       avl_set_child(X, -sign, B);
+       }
+}
+
+/* See description in avl_handle_subtree_growth()  */
+static AVL_INLINE struct avl_tree_node *
+avl_do_double_rotate(struct avl_tree_node * const B,
+                    struct avl_tree_node * const A, const int sign)
+{
+
+       struct avl_tree_node * const E = avl_get_child(B, -sign);
+       const int e = avl_get_balance_factor(E);
+
+       avl_rotate(B, +sign);
+       avl_rotate(A, -sign);
+       avl_set_balance_factor(A, ((sign * e <= 0) ? 0 : -e));
+       avl_set_balance_factor(B, ((sign * e >= 0) ? 0 : -e));
+       avl_set_balance_factor(E, 0);
+
+       return E;
+}
+
+static AVL_INLINE bool
+avl_handle_subtree_growth(struct avl_tree_node * const node,
+                         struct avl_tree_node * const parent,
+                         const int sign)
+{
+       /* Height of @node subtree of @parent has increased by 1.
+        * Adjust @parent's balance factor and check whether rotations need to
+        * be done.  */
+
+       const int old_balance_factor = avl_get_balance_factor(parent);
+       const int new_balance_factor = old_balance_factor + sign;
+
+       if (old_balance_factor == 0) {
+               /* Parent has increased in height but is still sufficiently
+                * balanced.  Continue up the tree.  */
+               avl_adjust_balance_factor(parent, sign);
+               return false;
+       }
+
+       if (new_balance_factor == 0) {
+               /* Parent is balanced; nothing more to to.  */
+               avl_adjust_balance_factor(parent, sign);
+               return true;
+       }
+
+       /* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
+        * The other case is symmetric --- that is, the rotations done are the
+        * the mirror images, all the balance factors are inverted, and left and
+        * right pointers are otherwise reversed.  */
+
+       /* Parent is left-heavy (balance_factor == -2).  */
+
+       if (sign * avl_get_balance_factor(node) > 0) {
+
+               /* Child node (@node, also B below) is also left-heavy.
+                * It must have balance_factor == -1.
+                * Do a clockwise ("right") rotation rooted at
+                * @parent (A below):
+                *
+                *           A              B
+                *          / \           /   \
+                *         B   C   =>    D     A
+                *        / \           / \   / \
+                *       D   E         F   G E   C
+                *      / \
+                *     F   G
+                *
+                * Before the rotation:
+                *      balance(A) = -2
+                *      balance(B) = -1
+                * Let x = height(C).  Then:
+                *      height(B) = x + 2
+                *      height(D) = x + 1
+                *      height(E) = x
+                *      max(height(F), height(G)) = x.
+                *
+                * After the rotation:
+                *      height(D) = max(height(F), height(G)) + 1
+                *                = x + 1
+                *      height(A) = max(height(E), height(C)) + 1
+                *                = max(x, x) + 1 = x + 1
+                *      balance(B) = 0
+                *      balance(A) = 0
+                */
+
+               avl_rotate(parent, -sign);
+
+               avl_set_balance_factor(node, 0);   /* B */
+               avl_set_balance_factor(parent, 0); /* A */
+       } else {
+               /* Child node (@node, also B below) is right-heavy.
+                * It must have balance_factor == +1.
+                * Do a counterclockwise ("left") rotation rooted at child node
+                * (B below), then a clockwise ("right") rotation rooted at
+                * parent node (A below).
+                *
+                *           A             A           E
+                *          / \           / \        /   \
+                *         B   C   =>    E   C  =>  B     A
+                *        / \           / \        / \   / \
+                *       D   E         B   G      D   F G   C
+                *          / \       / \
+                *         F   G     D   F
+                *
+                * Before the rotation:
+                *      balance(A) = -2
+                *      balance(B) = +1
+                * Let x = height(C).  Then:
+                *      height(A)
+                *      height(B) = x + 2
+                *      height(E) = x + 1
+                *      height(D) = x
+                *      max(height(F), height(G)) = x
+                *
+                * After both rotations:
+                *      height(A) = max(height(G), height(C)) + 1
+                *                = x + 1
+                *      balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
+                *      height(B) = max(height(D), height(F)) + 1
+                *                = x + 1
+                *      balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
+                *
+                *      height(E) = x + 2
+                *      balance(E) = 0
+                */
+               avl_do_double_rotate(node, parent, sign);
+       }
+       return true;
+}
+
+/* Rebalance the tree after insertion of the specified node.  */
+void
+avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
+                               struct avl_tree_node *inserted)
+{
+       struct avl_tree_node *node, *parent;
+       bool done = false;
+
+       inserted->left = NULL;
+       inserted->right = NULL;
+
+       for (node = inserted, parent = avl_get_parent(node);
+            parent && !done;
+            node = parent, parent = avl_get_parent(parent))
+       {
+               /* Height of @node subtree has increased by 1  */
+
+               if (node == parent->left)
+                       done = avl_handle_subtree_growth(node, parent, -1);
+               else
+                       done = avl_handle_subtree_growth(node, parent, +1);
+       }
+       /* Due to rotations, *root_ptr may no longer be the root of the tree  */
+       while (avl_get_parent(*root_ptr))
+               *root_ptr = avl_get_parent(*root_ptr);
+}
+
+static AVL_INLINE struct avl_tree_node *
+avl_handle_subtree_shrink(struct avl_tree_node *parent,
+                         const int sign,
+                         bool * const left_deleted_ret)
+{
+       struct avl_tree_node *node;
+
+       const int old_balance_factor = avl_get_balance_factor(parent);
+       const int new_balance_factor = old_balance_factor + sign;
+
+       if (old_balance_factor == 0) {
+               /* Prior to the deletion, the subtree rooted at
+                * @parent was perfectly balanced.  It's now
+                * unbalanced by 1, but that's okay and its height
+                * hasn't changed.  Nothing more to do.  */
+               avl_adjust_balance_factor(parent, sign);
+               return NULL;
+       } else if (new_balance_factor == 0) {
+               /* The subtree rooted at @parent is now perfectly
+                * balanced, whereas before the deletion it was
+                * unbalanced by 1.  Its height must have decreased
+                * by 1.  No rotation is needed at this location,
+                * but continue up the tree.  */
+               avl_adjust_balance_factor(parent, sign);
+               node = parent;
+       } else {
+               /* The subtree rooted at @parent is now significantly
+                * unbalanced (by 2 in some direction).  */
+               node = avl_get_child(parent, sign);
+
+               /* The rotations below are similar to those done during
+                * insertion.  The only new case is the one where the
+                * child node has a balance factor of 0.  */
+
+               if (sign * avl_get_balance_factor(node) >= 0) {
+                       avl_rotate(parent, -sign);
+
+                       if (avl_get_balance_factor(node) == 0) {
+                               avl_set_balance_factor(node,   -sign);
+                               avl_set_balance_factor(parent, +sign);
+                               /* Height is unchanged; nothing more to do.  */
+                               return NULL;
+                       } else {
+                               avl_set_balance_factor(node, 0);
+                               avl_set_balance_factor(parent, 0);
+                       }
+               } else {
+                       node = avl_do_double_rotate(node, parent, sign);
+               }
+       }
+       parent = avl_get_parent(node);
+       if (parent)
+               *left_deleted_ret = (node == parent->left);
+       return parent;
+}
+
+/* Swaps node X, which must have 2 children, with its in-order successor.
+ * This adjusts all the necessary pointers and balance factors.  */
+static AVL_INLINE void
+avl_tree_swap_with_successor(struct avl_tree_node * const X)
+{
+       struct avl_tree_node * const Y          = avl_tree_next_in_order(X);
+       struct avl_tree_node * const Y_right    = Y->right;
+       struct avl_tree_node * const Y_parent   = avl_get_parent(Y);
+       struct avl_tree_node * const X_parent   = avl_get_parent(X);
+       const int Y_balance_factor              = avl_get_balance_factor(Y);
+
+       /* Set child pointer to Y when placed in X's position  */
+       if (X_parent) {
+               if (X == X_parent->left)
+                       X_parent->left = Y;
+               else
+                       X_parent->right = Y;
+       }
+
+       /* Set child pointer to X when placed in Y's position.  Also set
+        * the right pointer of Y (it may point to X if the nodes are
+        * adjacent)  */
+       if (Y_parent != X) {
+               if (Y == Y_parent->left)
+                       Y_parent->left = X;
+               else
+                       Y_parent->right = X;
+               Y->right = X->right;
+       } else {
+               Y->right = X;
+       }
+
+       /* Set parent pointer to X when placed in Y's position  */
+       if (Y_right)
+               avl_set_parent(Y_right, X);
+
+       /* Note: Y has no left child since it is the in-order sucessor of X.  */
+
+       /* Set parent pointers to Y when placed in X's position, as well as X's
+        * parent when placed in Y's position.  */
+       avl_set_parent(X->left, Y);
+       if (X->right != Y) {
+               avl_set_parent(X->right, Y);
+               avl_set_parent(X, Y_parent);
+       } else {
+               avl_set_parent(X, Y);
+       }
+
+       avl_set_parent(Y, X_parent);
+       Y->left = X->left;
+       avl_set_balance_factor(Y, avl_get_balance_factor(X));
+
+       X->left = NULL;
+       X->right = Y_right;
+       avl_set_balance_factor(X, Y_balance_factor);
+}
+
+/* Removes the specified @node from the AVL tree whose root is pointed to by
+ * @root_ptr.
+ *
+ * This *only* unlinks the node and rebalances the tree; it does not free any
+ * memory or anything.  */
+void
+avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
+{
+       struct avl_tree_node *child, *parent;
+       bool left_deleted;
+
+       if (node->left && node->right)
+               avl_tree_swap_with_successor(node);
+
+       /* Unlink @node  */
+       child = node->left ? node->left : node->right;
+       parent = avl_get_parent(node);
+       if (parent) {
+               if (node == parent->left) {
+                       parent->left = child;
+                       left_deleted = true;
+               } else {
+                       parent->right = child;
+                       left_deleted = false;
+               }
+       } else {
+               *root_ptr = child;
+       }
+       if (child)
+               avl_set_parent(child, parent);
+
+       /* Rebalance the tree  */
+       while (parent) {
+               if (left_deleted)
+                       parent = avl_handle_subtree_shrink(parent, +1, &left_deleted);
+               else
+                       parent = avl_handle_subtree_shrink(parent, -1, &left_deleted);
+       }
+
+       /* Due to rotations, *root_ptr may no longer point to the root of the
+        * tree.  Fix it.  */
+       if (*root_ptr)
+               while (avl_get_parent(*root_ptr))
+                       *root_ptr = avl_get_parent(*root_ptr);
+}