Replace red-black tree code with AVL tree code
authorEric Biggers <ebiggers3@gmail.com>
Sat, 22 Mar 2014 03:38:54 +0000 (22:38 -0500)
committerEric Biggers <ebiggers3@gmail.com>
Sun, 23 Mar 2014 06:57:33 +0000 (01:57 -0500)
Makefile.am
include/wimlib/avl_tree.h [new file with mode: 0644]
include/wimlib/rbtree.h [deleted file]
src/avl_tree.c [new file with mode: 0644]
src/rbtree.c [deleted file]

index f5ac406..783821e 100644 (file)
@@ -19,6 +19,7 @@ libwim_la_LDFLAGS = -version-info 14:0:5 $(WINDOWS_LDFLAGS)
 
 libwim_la_SOURCES =            \
        src/add_image.c         \
+       src/avl_tree.c          \
        src/capture_common.c    \
        src/compat.c            \
        src/compress.c          \
@@ -58,7 +59,6 @@ libwim_la_SOURCES =           \
        src/pathlist.c          \
        src/paths.c             \
        src/resource.c          \
-       src/rbtree.c            \
        src/reference.c         \
        src/security.c          \
        src/sha1.c              \
@@ -77,6 +77,7 @@ libwim_la_SOURCES =           \
        src/xpress-decompress.c \
        include/wimlib/apply.h          \
        include/wimlib/assert.h         \
+       include/wimlib/avl_tree.h       \
        include/wimlib/callback.h       \
        include/wimlib/capture.h        \
        include/wimlib/case.h           \
@@ -107,7 +108,6 @@ libwim_la_SOURCES =         \
        include/wimlib/metadata.h       \
        include/wimlib/pathlist.h       \
        include/wimlib/paths.h          \
-       include/wimlib/rbtree.h         \
        include/wimlib/reparse.h        \
        include/wimlib/resource.h       \
        include/wimlib/security.h       \
diff --git a/include/wimlib/avl_tree.h b/include/wimlib/avl_tree.h
new file mode 100644 (file)
index 0000000..00fdd3e
--- /dev/null
@@ -0,0 +1,208 @@
+/*
+ * avl_tree.h
+ *
+ * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
+ * tree), header file.
+ *
+ * Author:  Eric Biggers
+ * Year:    2014
+ *
+ * This file is placed into the public domain.  You can do whatever you want
+ * with it.
+ */
+
+#ifndef _AVL_TREE_H_
+#define _AVL_TREE_H_
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <inttypes.h>          /* for uintptr_t */
+#include <wimlib/compiler.h>   /* for _always_inline_attribute. You can
+                                  instead define it below if you want.  */
+
+#define AVL_INLINE _always_inline_attribute
+
+/* Node in an AVL tree.  Embed this in some other data structure.  */
+struct avl_tree_node {
+
+       /* Pointer to left child or NULL  */
+       struct avl_tree_node *left;
+
+       /* Pointer to right child or NULL  */
+       struct avl_tree_node *right;
+
+       /* Pointer to parent combined with the balance factor.  This saves 4 or
+        * 8 bytes of memory depending on the CPU architecture.
+        *
+        * Low 2 bits:  One greater than the balance factor of this subtree,
+        * which is equal to height(right) - height(left).  The mapping is:
+        *
+        * 00 => -1
+        * 01 =>  0
+        * 10 => +1
+        * 11 => undefined
+        *
+        * The rest of the bits are the pointer to the parent node.  It must be
+        * 4-byte aligned, and it will be NULL if this is the root note and
+        * therefore has no parent.  */
+       uintptr_t parent_balance;
+};
+
+/* Cast an AVL tree node to the containing data structure.  */
+#define avl_tree_entry(entry, type, member) \
+       ((type*) ((char *)(entry) - offsetof(type, member)))
+
+/* Extracts the parent pointer from the specified AVL tree node.
+ * Returns the parent pointer, or NULL if @node is the root node.  */
+static AVL_INLINE struct avl_tree_node *
+avl_get_parent(const struct avl_tree_node *node)
+{
+       return (struct avl_tree_node *)(node->parent_balance & ~3);
+}
+
+/* Mark the node as unlinked from any tree.  */
+static AVL_INLINE void
+avl_tree_node_set_unlinked(struct avl_tree_node *node)
+{
+       node->parent_balance = (uintptr_t)node;
+}
+
+/* Returns true iff the specified node has been marked with
+ * avl_tree_node_set_unlinked() and has not subsequently been inserted into a
+ * tree.  */
+static AVL_INLINE bool
+avl_tree_node_is_unlinked(const struct avl_tree_node *node)
+{
+       return node->parent_balance == (uintptr_t)node;
+}
+
+/* Internal use only */
+extern void
+avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
+                               struct avl_tree_node *inserted);
+
+/*
+ * Look up an item in an AVL tree.
+ *
+ * @root
+ *     Pointer to the root of the AVL tree.  (This can be NULL --- that just
+ *     means the tree is empty.)
+ *
+ * @cmp_ctx
+ *     First argument to pass to the comparison callback --- generally a
+ *     pointer to an object equal to the one being searched for.
+ *
+ * @cmp
+ *     Comparison callback.  Must return < 0, 0, or > 0 if the first argument
+ *     is less than, equal to, or greater than the second argument,
+ *     respectively.  The first argument will be @cmp_ctx and the second
+ *     argument will be a pointer to the AVL tree node contained in the item
+ *     inserted into the AVL tree.
+ *
+ * Returns a pointer to the AVL tree node embedded in the resulting item, or
+ * NULL if the item was not found.
+ */
+static AVL_INLINE struct avl_tree_node *
+avl_tree_lookup(const struct avl_tree_node *root,
+               const void *cmp_ctx,
+               int (*cmp)(const void *, const struct avl_tree_node *))
+{
+       const struct avl_tree_node *cur = root;
+
+       while (cur) {
+               int res = (*cmp)(cmp_ctx, cur);
+               if (res < 0)
+                       cur = cur->left;
+               else if (res > 0)
+                       cur = cur->right;
+               else
+                       break;
+       }
+       return (struct avl_tree_node*)cur;
+}
+
+/* Same as avl_tree_lookup(), just uses a more specific type for the comparison
+ * function.  Specifically, the item being searched for is expected to be in the
+ * same format as those already in the tree, with an embedded 'struct
+ * avl_tree_node'.  */
+static AVL_INLINE struct avl_tree_node *
+avl_tree_lookup_node(const struct avl_tree_node *root,
+                    const struct avl_tree_node *node,
+                    int (*cmp)(const struct avl_tree_node *,
+                               const struct avl_tree_node *))
+{
+       return avl_tree_lookup(root,
+                              (const void *)node,
+                              (int (*) (const void *,
+                                        const struct avl_tree_node *))cmp);
+}
+
+/*
+ * Insert an item into an AVL tree.
+ *
+ * @root_ptr
+ *     Pointer to the pointer to the root of the AVL tree.  (Indirection is
+ *     needed because the root node may change.)  Initialize *root_ptr to NULL
+ *     for an empty tree.
+ *
+ * @item
+ *     Pointer to the `struct avl_tree_node' embedded in the item to insert.
+ *     No members in it need be pre-initialized, although members in the
+ *     containing structure should be pre-initialized so that @cmp can use them
+ *     in comparisons.
+ *
+ * @cmp
+ *     Comparison callback.  Must return < 0, 0, or > 0 if the first argument
+ *     is less than, equal to, or greater than the second argument,
+ *     respectively.  The first argument will be @item and the second
+ *     argument will be a pointer to an AVL tree node embedded in some
+ *     previously-inserted item that @item is being compared with.
+ *
+ * Returns NULL if the item was successfully inserted, otherwise the node of a
+ * previously-inserted item which compared equal to @item and prevented the new
+ * insertion of @item.
+ */
+static AVL_INLINE struct avl_tree_node *
+avl_tree_insert(struct avl_tree_node **root_ptr,
+               struct avl_tree_node *item,
+               int (*cmp)(const struct avl_tree_node *,
+                          const struct avl_tree_node *))
+{
+       struct avl_tree_node **cur_ptr = root_ptr, *cur = NULL;
+       int res;
+
+       while (*cur_ptr) {
+               cur = *cur_ptr;
+               res = (*cmp)(item, cur);
+               if (res < 0)
+                       cur_ptr = &cur->left;
+               else if (res > 0)
+                       cur_ptr = &cur->right;
+               else
+                       return cur;
+       }
+       *cur_ptr = item;
+       item->parent_balance = (uintptr_t)cur | 1;
+       avl_tree_rebalance_after_insert(root_ptr, item);
+       return NULL;
+}
+
+extern void
+avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node);
+
+/* Nonrecursive AVL tree traversal functions  */
+
+extern struct avl_tree_node *
+avl_tree_first_in_order(const struct avl_tree_node *root);
+
+extern struct avl_tree_node *
+avl_tree_next_in_order(const struct avl_tree_node *prev);
+
+extern struct avl_tree_node *
+avl_tree_first_in_postorder(const struct avl_tree_node *root);
+
+extern struct avl_tree_node *
+avl_tree_next_in_postorder(const struct avl_tree_node *prev,
+                          const struct avl_tree_node *prev_parent);
+
+#endif /* _AVL_TREE_H_ */
diff --git a/include/wimlib/rbtree.h b/include/wimlib/rbtree.h
deleted file mode 100644 (file)
index 3ae97a5..0000000
+++ /dev/null
@@ -1,93 +0,0 @@
-/*
- *  Red Black Trees
- *  Copyright (C) 1999  Andrea Arcangeli <andrea@suse.de>
- *
- *  This program is free software; you can redistribute it and/or modify
- *  it under the terms of the GNU General Public License as published by
- *  the Free Software Foundation; either version 2 of the License, or
- *  (at your option) any later version.
- *
- *  This program is distributed in the hope that it will be useful,
- *  but WITHOUT ANY WARRANTY; without even the implied warranty of
- *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *  GNU General Public License for more details.
- *
- *  You should have received a copy of the GNU General Public License
- *  along with this program; if not, write to the Free Software
- *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
- */
-
-#ifndef        _RBTREE_H
-#define        _RBTREE_H
-
-#include <stdint.h>
-#include <stddef.h>
-#include <stdbool.h>
-
-struct rb_node {
-       uintptr_t __rb_parent_color;
-       struct rb_node *rb_right;
-       struct rb_node *rb_left;
-};
-
-struct rb_root {
-       struct rb_node *rb_node;
-};
-
-#define        rb_entry(ptr, type, member) container_of(ptr, type, member)
-
-#define RB_ROOT        (struct rb_root) { NULL, }
-
-static inline bool
-rb_empty_root(const struct rb_root *root)
-{
-       return root->rb_node == NULL;
-}
-
-static inline struct rb_node *
-rb_parent(const struct rb_node *node)
-{
-       return (struct rb_node *)(node->__rb_parent_color & ~1);
-}
-
-/* 'empty' nodes are nodes that are known not to be inserted in an rbtree */
-static inline bool
-rb_empty_node(const struct rb_node *node)
-{
-       return node->__rb_parent_color == (uintptr_t)node;
-}
-
-static inline void
-rb_clear_node(struct rb_node *node)
-{
-       node->__rb_parent_color = (uintptr_t)node;
-}
-
-extern void
-rb_insert_color(struct rb_node *node, struct rb_root *root);
-
-extern void
-rb_erase(struct rb_node *node, struct rb_root *root);
-
-extern struct rb_node *
-rb_first(const struct rb_root *root);
-
-extern struct rb_node *
-rb_first_postorder(const struct rb_root *root);
-
-extern struct rb_node *
-rb_next(const struct rb_node *node);
-
-extern struct rb_node *
-rb_next_postorder(const struct rb_node *node, const struct rb_node *parent);
-
-static inline void
-rb_link_node(struct rb_node *node, struct rb_node *parent,
-            struct rb_node **rb_link)
-{
-       node->__rb_parent_color = (uintptr_t)parent;
-       node->rb_left = node->rb_right = NULL;
-       *rb_link = node;
-}
-
-#endif /* _RBTREE_H */
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);
+}
diff --git a/src/rbtree.c b/src/rbtree.c
deleted file mode 100644 (file)
index 5adbada..0000000
+++ /dev/null
@@ -1,557 +0,0 @@
-/*
- *  Red Black Trees
- *  Copyright (C) 1999  Andrea Arcangeli <andrea@suse.de>
- *  Copyright (C) 2002  David Woodhouse <dwmw2@infradead.org>
- *  Copyright (C) 2012  Michel Lespinasse <walken@google.com>
- *
- * This program is free software; you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation; either version 2 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program; if not, write to the Free Software
- * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
-*/
-
-#ifdef HAVE_CONFIG_H
-#  include "config.h"
-#endif
-
-#include "wimlib/rbtree.h"
-#include <stdbool.h>
-
-/*
- * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
- *
- *  1) A node is either red or black
- *  2) The root is black
- *  3) All leaves (NULL) are black
- *  4) Both children of every red node are black
- *  5) Every simple path from root to leaves contains the same number
- *     of black nodes.
- *
- *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
- *  consecutive red nodes in a path and every red node is therefore followed by
- *  a black. So if B is the number of black nodes on every simple path (as per
- *  5), then the longest possible path due to 4 is 2B.
- *
- *  We shall indicate color with case, where black nodes are uppercase and red
- *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
- *  parentheses and have some accompanying text comment.
- */
-
-#define        RB_RED          0
-#define        RB_BLACK        1
-
-#define __rb_parent(pc)    ((struct rb_node *)(pc & ~1))
-
-#define __rb_color(pc)     ((pc) & 1)
-#define __rb_is_black(pc)  __rb_color(pc)
-#define __rb_is_red(pc)    (!__rb_color(pc))
-#define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
-#define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
-#define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)
-
-static inline void
-rb_set_parent(struct rb_node *rb, struct rb_node *p)
-{
-       rb->__rb_parent_color = rb_color(rb) | (uintptr_t)p;
-}
-
-static inline void
-rb_set_parent_color(struct rb_node *rb, struct rb_node *p, int color)
-{
-       rb->__rb_parent_color = (uintptr_t)p | color;
-}
-
-static inline void
-rb_set_black(struct rb_node *rb)
-{
-       rb->__rb_parent_color |= RB_BLACK;
-}
-
-static inline struct rb_node *
-rb_red_parent(struct rb_node *red)
-{
-       return (struct rb_node *)red->__rb_parent_color;
-}
-
-static inline void
-rb_change_child(struct rb_node *old, struct rb_node *new,
-               struct rb_node *parent, struct rb_root *root)
-{
-       if (parent) {
-               if (parent->rb_left == old)
-                       parent->rb_left = new;
-               else
-                       parent->rb_right = new;
-       } else
-               root->rb_node = new;
-}
-
-/*
- * Helper function for rotations:
- * - old's parent and color get assigned to new
- * - old gets assigned new as a parent and 'color' as a color.
- */
-static inline void
-rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
-                       struct rb_root *root, int color)
-{
-       struct rb_node *parent = rb_parent(old);
-       new->__rb_parent_color = old->__rb_parent_color;
-       rb_set_parent_color(old, new, color);
-       rb_change_child(old, new, parent, root);
-}
-
-static void
-rb_erase_color(struct rb_node *parent, struct rb_root *root)
-{
-       struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
-
-       for (;;) {
-               /*
-                * Loop invariants:
-                * - node is black (or NULL on first iteration)
-                * - node is not the root (parent is not NULL)
-                * - All leaf paths going through parent and node have a
-                *   black node count that is 1 lower than other leaf paths.
-                */
-               sibling = parent->rb_right;
-               if (node != sibling) {  /* node == parent->rb_left */
-                       if (rb_is_red(sibling)) {
-                               /*
-                                * Case 1 - left rotate at parent
-                                *
-                                *     P               S
-                                *    / \             / \
-                                *   N   s    -->    p   Sr
-                                *      / \         / \
-                                *     Sl  Sr      N   Sl
-                                */
-                               parent->rb_right = tmp1 = sibling->rb_left;
-                               sibling->rb_left = parent;
-                               rb_set_parent_color(tmp1, parent, RB_BLACK);
-                               rb_rotate_set_parents(parent, sibling, root, RB_RED);
-                               sibling = tmp1;
-                       }
-                       tmp1 = sibling->rb_right;
-                       if (!tmp1 || rb_is_black(tmp1)) {
-                               tmp2 = sibling->rb_left;
-                               if (!tmp2 || rb_is_black(tmp2)) {
-                                       /*
-                                        * Case 2 - sibling color flip
-                                        * (p could be either color here)
-                                        *
-                                        *    (p)           (p)
-                                        *    / \           / \
-                                        *   N   S    -->  N   s
-                                        *      / \           / \
-                                        *     Sl  Sr        Sl  Sr
-                                        *
-                                        * This leaves us violating 5) which
-                                        * can be fixed by flipping p to black
-                                        * if it was red, or by recursing at p.
-                                        * p is red when coming from Case 1.
-                                        */
-                                       rb_set_parent_color(sibling, parent,
-                                                           RB_RED);
-                                       if (rb_is_red(parent))
-                                               rb_set_black(parent);
-                                       else {
-                                               node = parent;
-                                               parent = rb_parent(node);
-                                               if (parent)
-                                                       continue;
-                                       }
-                                       break;
-                               }
-                               /*
-                                * Case 3 - right rotate at sibling
-                                * (p could be either color here)
-                                *
-                                *   (p)           (p)
-                                *   / \           / \
-                                *  N   S    -->  N   Sl
-                                *     / \             \
-                                *    sl  Sr            s
-                                *                       \
-                                *                        Sr
-                                */
-                               sibling->rb_left = tmp1 = tmp2->rb_right;
-                               tmp2->rb_right = sibling;
-                               parent->rb_right = tmp2;
-                               if (tmp1)
-                                       rb_set_parent_color(tmp1, sibling,
-                                                           RB_BLACK);
-                               tmp1 = sibling;
-                               sibling = tmp2;
-                       }
-                       /*
-                        * Case 4 - left rotate at parent + color flips
-                        * (p and sl could be either color here.
-                        *  After rotation, p becomes black, s acquires
-                        *  p's color, and sl keeps its color)
-                        *
-                        *      (p)             (s)
-                        *      / \             / \
-                        *     N   S     -->   P   Sr
-                        *        / \         / \
-                        *      (sl) sr      N  (sl)
-                        */
-                       parent->rb_right = tmp2 = sibling->rb_left;
-                       sibling->rb_left = parent;
-                       rb_set_parent_color(tmp1, sibling, RB_BLACK);
-                       if (tmp2)
-                               rb_set_parent(tmp2, parent);
-                       rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
-                       break;
-               } else {
-                       sibling = parent->rb_left;
-                       if (rb_is_red(sibling)) {
-                               /* Case 1 - right rotate at parent */
-                               parent->rb_left = tmp1 = sibling->rb_right;
-                               sibling->rb_right = parent;
-                               rb_set_parent_color(tmp1, parent, RB_BLACK);
-                               rb_rotate_set_parents(parent, sibling, root,
-                                                     RB_RED);
-                               sibling = tmp1;
-                       }
-                       tmp1 = sibling->rb_left;
-                       if (!tmp1 || rb_is_black(tmp1)) {
-                               tmp2 = sibling->rb_right;
-                               if (!tmp2 || rb_is_black(tmp2)) {
-                                       /* Case 2 - sibling color flip */
-                                       rb_set_parent_color(sibling, parent,
-                                                           RB_RED);
-                                       if (rb_is_red(parent))
-                                               rb_set_black(parent);
-                                       else {
-                                               node = parent;
-                                               parent = rb_parent(node);
-                                               if (parent)
-                                                       continue;
-                                       }
-                                       break;
-                               }
-                               /* Case 3 - right rotate at sibling */
-                               sibling->rb_right = tmp1 = tmp2->rb_left;
-                               tmp2->rb_left = sibling;
-                               parent->rb_left = tmp2;
-                               if (tmp1)
-                                       rb_set_parent_color(tmp1, sibling,
-                                                           RB_BLACK);
-                               tmp1 = sibling;
-                               sibling = tmp2;
-                       }
-                       /* Case 4 - left rotate at parent + color flips */
-                       parent->rb_left = tmp2 = sibling->rb_right;
-                       sibling->rb_right = parent;
-                       rb_set_parent_color(tmp1, sibling, RB_BLACK);
-                       if (tmp2)
-                               rb_set_parent(tmp2, parent);
-                       rb_rotate_set_parents(parent, sibling, root, RB_BLACK);
-                       break;
-               }
-       }
-}
-
-void
-rb_erase(struct rb_node *node, struct rb_root *root)
-{
-       struct rb_node *child = node->rb_right, *tmp = node->rb_left;
-       struct rb_node *parent, *rebalance;
-       uintptr_t pc;
-
-       if (!tmp) {
-               /*
-                * Case 1: node to erase has no more than 1 child (easy!)
-                *
-                * Note that if there is one child it must be red due to 5)
-                * and node must be black due to 4). We adjust colors locally
-                * so as to bypass __rb_erase_color() later on.
-                */
-               pc = node->__rb_parent_color;
-               parent = __rb_parent(pc);
-               rb_change_child(node, child, parent, root);
-               if (child) {
-                       child->__rb_parent_color = pc;
-                       rebalance = NULL;
-               } else
-                       rebalance = __rb_is_black(pc) ? parent : NULL;
-               tmp = parent;
-       } else if (!child) {
-               /* Still case 1, but this time the child is node->rb_left */
-               tmp->__rb_parent_color = pc = node->__rb_parent_color;
-               parent = __rb_parent(pc);
-               rb_change_child(node, tmp, parent, root);
-               rebalance = NULL;
-               tmp = parent;
-       } else {
-               struct rb_node *successor = child, *child2;
-               tmp = child->rb_left;
-               if (!tmp) {
-                       /*
-                        * Case 2: node's successor is its right child
-                        *
-                        *    (n)          (s)
-                        *    / \          / \
-                        *  (x) (s)  ->  (x) (c)
-                        *        \
-                        *        (c)
-                        */
-                       parent = successor;
-                       child2 = successor->rb_right;
-               } else {
-                       /*
-                        * Case 3: node's successor is leftmost under
-                        * node's right child subtree
-                        *
-                        *    (n)          (s)
-                        *    / \          / \
-                        *  (x) (y)  ->  (x) (y)
-                        *      /            /
-                        *    (p)          (p)
-                        *    /            /
-                        *  (s)          (c)
-                        *    \
-                        *    (c)
-                        */
-                       do {
-                               parent = successor;
-                               successor = tmp;
-                               tmp = tmp->rb_left;
-                       } while (tmp);
-                       parent->rb_left = child2 = successor->rb_right;
-                       successor->rb_right = child;
-                       rb_set_parent(child, successor);
-               }
-
-               successor->rb_left = tmp = node->rb_left;
-               rb_set_parent(tmp, successor);
-
-               pc = node->__rb_parent_color;
-               tmp = __rb_parent(pc);
-               rb_change_child(node, successor, tmp, root);
-               if (child2) {
-                       successor->__rb_parent_color = pc;
-                       rb_set_parent_color(child2, parent, RB_BLACK);
-                       rebalance = NULL;
-               } else {
-                       uintptr_t pc2 = successor->__rb_parent_color;
-                       successor->__rb_parent_color = pc;
-                       rebalance = __rb_is_black(pc2) ? parent : NULL;
-               }
-               tmp = successor;
-       }
-
-       if (rebalance)
-               rb_erase_color(rebalance, root);
-}
-
-void
-rb_insert_color(struct rb_node *node, struct rb_root *root)
-{
-       struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
-
-       while (true) {
-               /*
-                * Loop invariant: node is red
-                *
-                * If there is a black parent, we are done.
-                * Otherwise, take some corrective action as we don't
-                * want a red root or two consecutive red nodes.
-                */
-               if (!parent) {
-                       rb_set_parent_color(node, NULL, RB_BLACK);
-                       break;
-               } else if (rb_is_black(parent))
-                       break;
-
-               gparent = rb_red_parent(parent);
-
-               tmp = gparent->rb_right;
-               if (parent != tmp) {    /* parent == gparent->rb_left */
-                       if (tmp && rb_is_red(tmp)) {
-                               /*
-                                * Case 1 - color flips
-                                *
-                                *       G            g
-                                *      / \          / \
-                                *     p   u  -->   P   U
-                                *    /            /
-                                *   n            N
-                                *
-                                * However, since g's parent might be red, and
-                                * 4) does not allow this, we need to recurse
-                                * at g.
-                                */
-                               rb_set_parent_color(tmp, gparent, RB_BLACK);
-                               rb_set_parent_color(parent, gparent, RB_BLACK);
-                               node = gparent;
-                               parent = rb_parent(node);
-                               rb_set_parent_color(node, parent, RB_RED);
-                               continue;
-                       }
-
-                       tmp = parent->rb_right;
-                       if (node == tmp) {
-                               /*
-                                * Case 2 - left rotate at parent
-                                *
-                                *      G             G
-                                *     / \           / \
-                                *    p   U  -->    n   U
-                                *     \           /
-                                *      n         p
-                                *
-                                * This still leaves us in violation of 4), the
-                                * continuation into Case 3 will fix that.
-                                */
-                               parent->rb_right = tmp = node->rb_left;
-                               node->rb_left = parent;
-                               if (tmp)
-                                       rb_set_parent_color(tmp, parent,
-                                                           RB_BLACK);
-                               rb_set_parent_color(parent, node, RB_RED);
-                               parent = node;
-                               tmp = node->rb_right;
-                       }
-
-                       /*
-                        * Case 3 - right rotate at gparent
-                        *
-                        *        G           P
-                        *       / \         / \
-                        *      p   U  -->  n   g
-                        *     /                 \
-                        *    n                   U
-                        */
-                       gparent->rb_left = tmp;  /* == parent->rb_right */
-                       parent->rb_right = gparent;
-                       if (tmp)
-                               rb_set_parent_color(tmp, gparent, RB_BLACK);
-                       rb_rotate_set_parents(gparent, parent, root, RB_RED);
-                       break;
-               } else {
-                       tmp = gparent->rb_left;
-                       if (tmp && rb_is_red(tmp)) {
-                               /* Case 1 - color flips */
-                               rb_set_parent_color(tmp, gparent, RB_BLACK);
-                               rb_set_parent_color(parent, gparent, RB_BLACK);
-                               node = gparent;
-                               parent = rb_parent(node);
-                               rb_set_parent_color(node, parent, RB_RED);
-                               continue;
-                       }
-
-                       tmp = parent->rb_left;
-                       if (node == tmp) {
-                               /* Case 2 - right rotate at parent */
-                               parent->rb_left = tmp = node->rb_right;
-                               node->rb_right = parent;
-                               if (tmp)
-                                       rb_set_parent_color(tmp, parent,
-                                                           RB_BLACK);
-                               rb_set_parent_color(parent, node, RB_RED);
-                               parent = node;
-                               tmp = node->rb_left;
-                       }
-
-                       /* Case 3 - left rotate at gparent */
-                       gparent->rb_right = tmp;  /* == parent->rb_left */
-                       parent->rb_left = gparent;
-                       if (tmp)
-                               rb_set_parent_color(tmp, gparent, RB_BLACK);
-                       rb_rotate_set_parents(gparent, parent, root, RB_RED);
-                       break;
-               }
-       }
-}
-
-static struct rb_node *
-rb_left_deepest_node(const struct rb_node *node)
-{
-       for (;;) {
-               if (node->rb_left)
-                       node = node->rb_left;
-               else if (node->rb_right)
-                       node = node->rb_right;
-               else
-                       return (struct rb_node *)node;
-       }
-}
-
-struct rb_node *
-rb_next_postorder(const struct rb_node *node, const struct rb_node *parent)
-{
-       /* If we're sitting on node, we've already seen our children */
-       if (parent && node == parent->rb_left && parent->rb_right) {
-               /* If we are the parent's left node, go to the parent's right
-                * node then all the way down to the left */
-               return rb_left_deepest_node(parent->rb_right);
-       } else
-               /* Otherwise we are the parent's right node, and the parent
-                * should be next */
-               return (struct rb_node *)parent;
-}
-
-struct rb_node *
-rb_first_postorder(const struct rb_root *root)
-{
-       if (!root->rb_node)
-               return NULL;
-
-       return rb_left_deepest_node(root->rb_node);
-}
-
-struct rb_node *
-rb_next(const struct rb_node *node)
-{
-       struct rb_node *parent;
-
-       /*
-        * If we have a right-hand child, go down and then left as far
-        * as we can.
-        */
-       if (node->rb_right) {
-               node = node->rb_right;
-               while (node->rb_left)
-                       node = node->rb_left;
-               return (struct rb_node *)node;
-       }
-
-       /*
-        * No right-hand children. Everything down and left is smaller than us,
-        * so any 'next' node must be in the general direction of our parent.
-        * Go up the tree; any time the ancestor is a right-hand child of its
-        * parent, keep going up. First time it's a left-hand child of its
-        * parent, said parent is our 'next' node.
-        */
-       while ((parent = rb_parent(node)) && node == parent->rb_right)
-               node = parent;
-
-       return parent;
-}
-
-/*
- * This function returns the first node (in sort order) of the tree.
- */
-struct rb_node *
-rb_first(const struct rb_root *root)
-{
-       struct rb_node  *n;
-
-       n = root->rb_node;
-       if (!n)
-               return NULL;
-       while (n->rb_left)
-               n = n->rb_left;
-       return n;
-}