]> wimlib.net Git - wimlib/blobdiff - src/avl_tree.c
avl_tree: Optimize first iteration of insertion rebalance loop
[wimlib] / src / avl_tree.c
index ab269d0e72d0c98c93f2f50d83b0dd6b83f0f62c..8b8e07bbe5d608b26975a47d60e443c044c5d033 100644 (file)
@@ -110,6 +110,14 @@ avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
                (node->parent_balance & 3) | (uintptr_t)parent;
 }
 
+/* Sets the parent and balance factor of the AVL tree @node.  */
+static AVL_INLINE void
+avl_set_parent_balance(struct avl_tree_node *node, struct avl_tree_node *parent,
+                      int balance_factor)
+{
+       node->parent_balance = (uintptr_t)parent | (balance_factor + 1);
+}
+
 /* 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
@@ -162,7 +170,8 @@ avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
  * This updates pointers but not balance factors!
  */
 static AVL_INLINE void
-avl_rotate(struct avl_tree_node * const A, const int sign)
+avl_rotate(struct avl_tree_node ** const root_ptr,
+          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);
@@ -184,38 +193,96 @@ avl_rotate(struct avl_tree_node * const A, const int sign)
                        avl_set_child(X, +sign, B);
                else
                        avl_set_child(X, -sign, B);
+       } else {
+               *root_ptr = B;
        }
 }
 
-/* See description in avl_handle_subtree_growth()  */
+/*
+ * Template for performing a double rotation ---
+ *
+ * Rotate counterclockwise (left) rooted at B, then clockwise (right) rooted at
+ * A (sign > 0):
+ *
+ *           A             A           E
+ *          / \           / \        /   \
+ *         B   C   =>    E   C  =>  B     A
+ *        / \           / \        / \   / \
+ *       D   E         B   G      D   F G   C
+ *          / \       / \
+ *         F   G     D   F
+ *
+ * Rotate clockwise (right) rooted at B, then counterclockwise (left) rooted at
+ * A (sign < 0):
+ *
+ *         A           A               E
+ *        / \         / \            /   \
+ *       C   B   =>  C   E    =>    A     B
+ *          / \         / \        / \   / \
+ *         E   D       G   B      C   G F   D
+ *        / \             / \
+ *       G   F           F   D
+ *
+ * Returns a pointer to E, the new root, and updates balance factors.  Except
+ * for that, this is equivalent to:
+ *     avl_rotate(root_ptr, B, -sign);
+ *     avl_rotate(root_ptr, A, +sign);
+ *
+ */
 static AVL_INLINE struct avl_tree_node *
-avl_do_double_rotate(struct avl_tree_node * const B,
+avl_do_double_rotate(struct avl_tree_node ** const root_ptr,
+                    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);
+       struct avl_tree_node * const E = avl_get_child(B, +sign);
+       struct avl_tree_node * const F = avl_get_child(E, -sign);
+       struct avl_tree_node * const G = avl_get_child(E, +sign);
+       struct avl_tree_node * const X = avl_get_parent(A);
        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);
+       avl_set_child(A, -sign, G);
+       avl_set_parent_balance(A, E, ((sign * e >= 0) ? 0 : -e));
+
+       avl_set_child(B, +sign, F);
+       avl_set_parent_balance(B, E, ((sign * e <= 0) ? 0 : -e));
+
+       avl_set_child(E, +sign, A);
+       avl_set_child(E, -sign, B);
+       avl_set_parent_balance(E, X, 0);
+
+       if (G)
+               avl_set_parent(G, A);
+
+       if (F)
+               avl_set_parent(F, B);
+
+       if (X) {
+               if (avl_get_child(X, +sign) == A)
+                       avl_set_child(X, +sign, E);
+               else
+                       avl_set_child(X, -sign, E);
+       } else {
+               *root_ptr = E;
+       }
 
        return E;
 }
 
 static AVL_INLINE bool
-avl_handle_subtree_growth(struct avl_tree_node * const node,
+avl_handle_subtree_growth(struct avl_tree_node ** const root_ptr,
+                         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.
+       /* The subtree rooted at @node, which is a child 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;
+       int old_balance_factor, new_balance_factor;
+
+       old_balance_factor = avl_get_balance_factor(parent);
 
        if (old_balance_factor == 0) {
                /* Parent has increased in height but is still sufficiently
@@ -224,6 +291,8 @@ avl_handle_subtree_growth(struct avl_tree_node * const node,
                return false;
        }
 
+       new_balance_factor = old_balance_factor + sign;
+
        if (new_balance_factor == 0) {
                /* Parent is balanced; nothing more to to.  */
                avl_adjust_balance_factor(parent, sign);
@@ -270,10 +339,9 @@ avl_handle_subtree_growth(struct avl_tree_node * const node,
                 *      balance(A) = 0
                 */
 
-               avl_rotate(parent, -sign);
-
-               avl_set_balance_factor(node, 0);   /* B */
-               avl_set_balance_factor(parent, 0); /* A */
+               avl_rotate(root_ptr, parent, -sign);
+               avl_adjust_balance_factor(parent, -sign); /* A */
+               avl_adjust_balance_factor(node, -sign);   /* B */
        } else {
                /* Child node (@node, also B below) is right-heavy.
                 * It must have balance_factor == +1.
@@ -293,7 +361,6 @@ avl_handle_subtree_growth(struct avl_tree_node * const node,
                 *      balance(A) = -2
                 *      balance(B) = +1
                 * Let x = height(C).  Then:
-                *      height(A)
                 *      height(B) = x + 2
                 *      height(E) = x + 1
                 *      height(D) = x
@@ -310,7 +377,7 @@ avl_handle_subtree_growth(struct avl_tree_node * const node,
                 *      height(E) = x + 2
                 *      balance(E) = 0
                 */
-               avl_do_double_rotate(node, parent, sign);
+               avl_do_double_rotate(root_ptr, node, parent, -sign);
        }
        return true;
 }
@@ -321,29 +388,52 @@ 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;
+       bool done;
 
        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  */
+       node = inserted;
+
+       /* Adjust balance factor of new node's parent.
+        * No rotation will need to be done at this level.  */
+
+       parent = avl_get_parent(node);
+       if (!parent)
+               return;
+
+       if (node == parent->left)
+               avl_adjust_balance_factor(parent, -1);
+       else
+               avl_adjust_balance_factor(parent, +1);
+
+       if (avl_get_balance_factor(parent) == 0)
+               /* @parent did not change in height.  Nothing more to do.  */
+               return;
+
+       /* The subtree rooted at @parent increased in height by 1.  */
+
+       do {
+               /* Adjust balance factor of next ancestor.  */
+
+               node = parent;
+               parent = avl_get_parent(node);
+               if (!parent)
+                       return;
 
+               /* The subtree rooted at @node has increased in height by 1.  */
                if (node == parent->left)
-                       done = avl_handle_subtree_growth(node, parent, -1);
+                       done = avl_handle_subtree_growth(root_ptr, 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);
+                       done = avl_handle_subtree_growth(root_ptr, node,
+                                                        parent, +1);
+       } while (!done);
 }
 
 static AVL_INLINE struct avl_tree_node *
-avl_handle_subtree_shrink(struct avl_tree_node *parent,
+avl_handle_subtree_shrink(struct avl_tree_node ** const root_ptr,
+                         struct avl_tree_node *parent,
                          const int sign,
                          bool * const left_deleted_ret)
 {
@@ -377,19 +467,57 @@ avl_handle_subtree_shrink(struct avl_tree_node *parent,
                 * child node has a balance factor of 0.  */
 
                if (sign * avl_get_balance_factor(node) >= 0) {
-                       avl_rotate(parent, -sign);
+                       avl_rotate(root_ptr, parent, -sign);
 
                        if (avl_get_balance_factor(node) == 0) {
-                               avl_set_balance_factor(node,   -sign);
-                               avl_set_balance_factor(parent, +sign);
+                               /* Child node (@node, also B below) is balanced.
+                                * 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) =  0
+                                * Let x = height(C).  Then:
+                                *      height(B) = x + 2
+                                *      height(D) = x + 1
+                                *      height(E) = x + 1
+                                *      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 + 1, x) + 1 = x + 2
+                                *      balance(A) = -1
+                                *      balance(B) = +1
+                                */
+
+                               /* A: -2 => -1 (sign < 0)
+                                * or +2 => +1 (sign > 0)
+                                * No change needed --- that's the same as
+                                * old_balance_factor.  */
+
+                               /* B: 0 => +1 (sign < 0)
+                                * or 0 => -1 (sign > 0)  */
+                               avl_adjust_balance_factor(node, -sign);
+
                                /* Height is unchanged; nothing more to do.  */
                                return NULL;
                        } else {
-                               avl_set_balance_factor(node, 0);
-                               avl_set_balance_factor(parent, 0);
+                               avl_adjust_balance_factor(parent, -sign);
+                               avl_adjust_balance_factor(node, -sign);
                        }
                } else {
-                       node = avl_do_double_rotate(node, parent, sign);
+                       node = avl_do_double_rotate(root_ptr, node,
+                                                   parent, -sign);
                }
        }
        parent = avl_get_parent(node);
@@ -398,105 +526,125 @@ avl_handle_subtree_shrink(struct avl_tree_node *parent,
        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)
+/* Swaps node X, which must have 2 children, with its in-order successor, then
+ * unlinks node X.  Returns the parent of X just before unlinking, without its
+ * balance factor having been updated to account for the unlink.  */
+static AVL_INLINE struct avl_tree_node *
+avl_tree_swap_with_successor(struct avl_tree_node **root_ptr,
+                            struct avl_tree_node *X,
+                            bool *left_deleted_ret)
 {
-       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;
-       }
+       struct avl_tree_node *Y, *P, *ret;
+
+       Y = X->right;
+       if (!Y->left) {
+               /*
+                *     P?           P?           P?
+                *     |            |            |
+                *     X            Y            Y
+                *    / \          / \          / \
+                *   A   Y    =>  A   X    =>  A   B?
+                *      / \          / \
+                *    (0)  B?      (0)  B?
+                *
+                * [ X removed, Y returned ]
+                */
+               ret = Y;
+               *left_deleted_ret = false;
+       } else {
+               struct avl_tree_node *Q;
+
+               do {
+                       Q = Y;
+                       Y = Y->left;
+               } while (Y->left);
+
+               /*
+                *     P?           P?           P?
+                *     |            |            |
+                *     X            Y            Y
+                *    / \          / \          / \
+                *   A   ...  =>  A  ...   =>  A  ...
+                *       |            |            |
+                *       Q            Q            Q
+                *      /            /            /
+                *     Y            X            B?
+                *    / \          / \
+                *  (0)  B?      (0)  B?
+                *
+                *
+                * [ X removed, Q returned ]
+                */
 
-       /* 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;
+               Q->left = Y->right;
+               if (Q->left)
+                       avl_set_parent(Q->left, Q);
                Y->right = X->right;
-       } else {
-               Y->right = X;
+               avl_set_parent(X->right, Y);
+               ret = Q;
+               *left_deleted_ret = true;
        }
 
-       /* 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.  */
+       Y->left = X->left;
        avl_set_parent(X->left, Y);
-       if (X->right != Y) {
-               avl_set_parent(X->right, Y);
-               avl_set_parent(X, Y_parent);
+
+       Y->parent_balance = X->parent_balance;
+       P = avl_get_parent(X);
+       if (P) {
+               if (P->left == X)
+                       P->left = Y;
+               else
+                       P->right = Y;
        } else {
-               avl_set_parent(X, Y);
+               *root_ptr = 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);
+       return ret;
 }
 
-/* Removes the specified @node from the AVL tree whose root is pointed to by
- * @root_ptr.
+/* Removes the specified @node from the AVL tree.  @root_ptr must point to the
+ * pointer to the root node of the tree; *root_ptr may change if the tree is
+ * rebalanced.
  *
- * This *only* unlinks the node and rebalances the tree; it does not free any
- * memory or anything.  */
+ * This *only* removes the node and rebalances the tree; it does not free
+ * memory, nor does it do the equivalent of avl_tree_node_set_unlinked().  */
 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;
+       if (node->left && node->right) {
+               parent = avl_tree_swap_with_successor(root_ptr, node,
+                                                     &left_deleted);
+       } else {
+               /* 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 {
-                       parent->right = child;
-                       left_deleted = false;
+                       *root_ptr = child;
                }
-       } else {
-               *root_ptr = child;
+               if (child)
+                       avl_set_parent(child, parent);
+               if (!parent)
+                       return;
        }
-       if (child)
-               avl_set_parent(child, parent);
 
        /* Rebalance the tree  */
-       while (parent) {
+       do {
                if (left_deleted)
-                       parent = avl_handle_subtree_shrink(parent, +1, &left_deleted);
+                       parent = avl_handle_subtree_shrink(root_ptr, 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);
+                       parent = avl_handle_subtree_shrink(root_ptr, parent,
+                                                          -1, &left_deleted);
+       } while (parent);
 }