wimlib.net Git - wimlib/blob - include/wimlib/avl_tree.h

   1 /*
   2  * avl_tree.h
   3  *
   4  * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
   5  * tree), header file.
   6  *
   7  * Author:  Eric Biggers
   8  * Year:    2014
   9  *
  10  * This file is placed into the public domain.  You can do whatever you want
  11  * with it.
  12  */
  13
  14 #ifndef _AVL_TREE_H_
  15 #define _AVL_TREE_H_
  16
  17 #include <stdbool.h>
  18 #include <stddef.h>
  19 #include <inttypes.h> /* for uintptr_t */
  20
  21 #ifdef __GNUC__
  22 #  define AVL_INLINE inline __attribute__((always_inline))
  23 #else
  24 #  define AVL_INLINE inline
  25 #  warning "AVL tree functions may not be inlined as intended"
  26 #endif
  27
  28 /* Node in an AVL tree.  Embed this in some other data structure.  */
  29 struct avl_tree_node {
  30
  31         /* Pointer to left child or NULL  */
  32         struct avl_tree_node *left;
  33
  34         /* Pointer to right child or NULL  */
  35         struct avl_tree_node *right;
  36
  37         /* Pointer to parent combined with the balance factor.  This saves 4 or
  38          * 8 bytes of memory depending on the CPU architecture.
  39          *
  40          * Low 2 bits:  One greater than the balance factor of this subtree,
  41          * which is equal to height(right) - height(left).  The mapping is:
  42          *
  43          * 00 => -1
  44          * 01 =>  0
  45          * 10 => +1
  46          * 11 => undefined
  47          *
  48          * The rest of the bits are the pointer to the parent node.  It must be
  49          * 4-byte aligned, and it will be NULL if this is the root node and
  50          * therefore has no parent.  */
  51         uintptr_t parent_balance;
  52 };
  53
  54 /* Cast an AVL tree node to the containing data structure.  */
  55 #define avl_tree_entry(entry, type, member) \
  56         ((type*) ((char *)(entry) - offsetof(type, member)))
  57
  58 /* Returns a pointer to the parent of the specified AVL tree node, or NULL if it
  59  * is already the root of the tree.  */
  60 static AVL_INLINE struct avl_tree_node *
  61 avl_get_parent(const struct avl_tree_node *node)
  62 {
  63         return (struct avl_tree_node *)(node->parent_balance & ~3);
  64 }
  65
  66 /* Marks the specified AVL tree node as unlinked from any tree.  */
  67 static AVL_INLINE void
  68 avl_tree_node_set_unlinked(struct avl_tree_node *node)
  69 {
  70         node->parent_balance = (uintptr_t)node;
  71 }
  72
  73 /* Returns true iff the specified AVL tree node has been marked with
  74  * avl_tree_node_set_unlinked() and has not subsequently been inserted into a
  75  * tree.  */
  76 static AVL_INLINE bool
  77 avl_tree_node_is_unlinked(const struct avl_tree_node *node)
  78 {
  79         return node->parent_balance == (uintptr_t)node;
  80 }
  81
  82 /* (Internal use only)  */
  83 extern void
  84 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
  85                                 struct avl_tree_node *inserted);
  86
  87 /*
  88  * Looks up an item in the specified AVL tree.
  89  *
  90  * @root
  91  *      Pointer to the root of the AVL tree.  (This can be NULL --- that just
  92  *      means the tree is empty.)
  93  *
  94  * @cmp_ctx
  95  *      First argument to pass to the comparison callback.  This generally
  96  *      should be a pointer to an object equal to the one being searched for.
  97  *
  98  * @cmp
  99  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
 100  *      is less than, equal to, or greater than the second argument,
 101  *      respectively.  The first argument will be @cmp_ctx and the second
 102  *      argument will be a pointer to the AVL tree node of an item in the tree.
 103  *
 104  * Returns a pointer to the AVL tree node of the resulting item, or NULL if the
 105  * item was not found.
 106  *
 107  * Example:
 108  *
 109  * struct int_wrapper {
 110  *      int data;
 111  *      struct avl_tree_node index_node;
 112  * };
 113  *
 114  * static int _avl_cmp_int_to_node(const void *intptr,
 115  *                                 const struct avl_tree_node *nodeptr)
 116  * {
 117  *      int n1 = *(const int *)intptr;
 118  *      int n2 = avl_tree_entry(nodeptr, struct int_wrapper, index_node)->data;
 119  *      if (n1 < n2)
 120  *              return -1;
 121  *      else if (n1 > n2)
 122  *              return 1;
 123  *      else
 124  *              return 0;
 125  * }
 126  *
 127  * bool contains_int(struct avl_tree_node *root, int n)
 128  * {
 129  *      struct avl_tree_node *result;
 130  *
 131  *      result = avl_tree_lookup(root, &n, _avl_cmp_int_to_node);
 132  *      return result ? true : false;
 133  * }
 134  */
 135 static AVL_INLINE struct avl_tree_node *
 136 avl_tree_lookup(const struct avl_tree_node *root,
 137                 const void *cmp_ctx,
 138                 int (*cmp)(const void *, const struct avl_tree_node *))
 139 {
 140         const struct avl_tree_node *cur = root;
 141
 142         while (cur) {
 143                 int res = (*cmp)(cmp_ctx, cur);
 144                 if (res < 0)
 145                         cur = cur->left;
 146                 else if (res > 0)
 147                         cur = cur->right;
 148                 else
 149                         break;
 150         }
 151         return (struct avl_tree_node*)cur;
 152 }
 153
 154 /* Same as avl_tree_lookup(), but uses a more specific type for the comparison
 155  * function.  Specifically, with this function the item being searched for is
 156  * expected to be in the same format as those already in the tree, with an
 157  * embedded 'struct avl_tree_node'.  */
 158 static AVL_INLINE struct avl_tree_node *
 159 avl_tree_lookup_node(const struct avl_tree_node *root,
 160                      const struct avl_tree_node *node,
 161                      int (*cmp)(const struct avl_tree_node *,
 162                                 const struct avl_tree_node *))
 163 {
 164         return avl_tree_lookup(root,
 165                                (const void *)node,
 166                                (int (*) (const void *,
 167                                          const struct avl_tree_node *))cmp);
 168 }
 169
 170 /*
 171  * Inserts an item into the specified AVL tree.
 172  *
 173  * @root_ptr
 174  *      Location of the AVL tree's root pointer.  Indirection is needed because
 175  *      the root node may change as a result of rotations caused by the
 176  *      insertion.  Initialize *root_ptr to NULL for an empty tree.
 177  *
 178  * @item
 179  *      Pointer to the `struct avl_tree_node' embedded in the item to insert.
 180  *      No members in it need be pre-initialized, although members in the
 181  *      containing structure should be pre-initialized so that @cmp can use them
 182  *      in comparisons.
 183  *
 184  * @cmp
 185  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
 186  *      is less than, equal to, or greater than the second argument,
 187  *      respectively.  The first argument will be @item and the second
 188  *      argument will be a pointer to an AVL tree node embedded in some
 189  *      previously-inserted item to which @item is being compared.
 190  *
 191  * If no item in the tree is comparatively equal (via @cmp) to @item, inserts
 192  * @item and returns NULL.  Otherwise does nothing and returns a pointer to the
 193  * AVL tree node embedded in the previously-inserted item which compared equal
 194  * to @item.
 195  *
 196  * Example:
 197  *
 198  * struct int_wrapper {
 199  *      int data;
 200  *      struct avl_tree_node index_node;
 201  * };
 202  *
 203  * #define GET_DATA(i) avl_tree_entry((i), struct int_wrapper, index_node)->data
 204  *
 205  * static int _avl_cmp_ints(const struct avl_tree_node *node1,
 206  *                          const struct avl_tree_node *node2)
 207  * {
 208  *      int n1 = GET_DATA(node1);
 209  *      int n2 = GET_DATA(node2);
 210  *      if (n1 < n2)
 211  *              return -1;
 212  *      else if (n1 > n2)
 213  *              return 1;
 214  *      else
 215  *              return 0;
 216  * }
 217  *
 218  * bool insert_int(struct avl_tree_node **root_ptr, int data)
 219  * {
 220  *      struct int_wrapper *i = malloc(sizeof(struct int_wrapper));
 221  *      i->data = data;
 222  *      if (avl_tree_insert(root_ptr, &i->index_node, _avl_cmp_ints)) {
 223  *              // Duplicate.
 224  *              free(i);
 225  *              return false;
 226  *      }
 227  *      return true;
 228  * }
 229  */
 230 static AVL_INLINE struct avl_tree_node *
 231 avl_tree_insert(struct avl_tree_node **root_ptr,
 232                 struct avl_tree_node *item,
 233                 int (*cmp)(const struct avl_tree_node *,
 234                            const struct avl_tree_node *))
 235 {
 236         struct avl_tree_node **cur_ptr = root_ptr, *cur = NULL;
 237         int res;
 238
 239         while (*cur_ptr) {
 240                 cur = *cur_ptr;
 241                 res = (*cmp)(item, cur);
 242                 if (res < 0)
 243                         cur_ptr = &cur->left;
 244                 else if (res > 0)
 245                         cur_ptr = &cur->right;
 246                 else
 247                         return cur;
 248         }
 249         *cur_ptr = item;
 250         item->parent_balance = (uintptr_t)cur | 1;
 251         avl_tree_rebalance_after_insert(root_ptr, item);
 252         return NULL;
 253 }
 254
 255 /* Removes an item from the specified AVL tree.
 256  * See implementation for details.  */
 257 extern void
 258 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node);
 259
 260 /* Nonrecursive AVL tree traversal functions  */
 261
 262 extern struct avl_tree_node *
 263 avl_tree_first_in_order(const struct avl_tree_node *root);
 264
 265 extern struct avl_tree_node *
 266 avl_tree_next_in_order(const struct avl_tree_node *prev);
 267
 268 extern struct avl_tree_node *
 269 avl_tree_first_in_postorder(const struct avl_tree_node *root);
 270
 271 extern struct avl_tree_node *
 272 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
 273                            const struct avl_tree_node *prev_parent);
 274
 275 /*
 276  * Iterate through the nodes in an AVL tree in sorted order.
 277  * You may not modify the tree during the iteration.
 278  *
 279  * @child_struct
 280  *      Variable that will receive a pointer to each struct inserted into the
 281  *      tree.
 282  * @root
 283  *      Root of the AVL tree.
 284  * @struct_name
 285  *      Type of *child_struct.
 286  * @struct_member
 287  *      Member of @struct_name type that is the AVL tree node.
 288  *
 289  * Example:
 290  *
 291  * struct int_wrapper {
 292  *      int data;
 293  *      struct avl_tree_node index_node;
 294  * };
 295  *
 296  * void print_ints(struct avl_tree_node *root)
 297  * {
 298  *      struct int_wrapper *i;
 299  *
 300  *      avl_tree_for_each_in_order(i, root, struct int_wrapper, index_node)
 301  *              printf("%d\n", i->data);
 302  * }
 303  */
 304 #define avl_tree_for_each_in_order(child_struct, root,                  \
 305                                    struct_name, struct_member)          \
 306         for (struct avl_tree_node *_cur =                               \
 307                 avl_tree_first_in_order(root);                          \
 308              _cur && ((child_struct) =                                  \
 309                       avl_tree_entry(_cur, struct_name,                 \
 310                                      struct_member), 1);                \
 311              _cur = avl_tree_next_in_order(_cur))
 312
 313 /*
 314  * Like avl_tree_for_each_in_order(), but iterates through the nodes in
 315  * postorder, so the current node may be deleted or freed.
 316  */
 317 #define avl_tree_for_each_in_postorder(child_struct, root,              \
 318                                        struct_name, struct_member)      \
 319         for (struct avl_tree_node *_cur =                               \
 320                 avl_tree_first_in_postorder(root), *_parent;            \
 321              _cur && ((child_struct) =                                  \
 322                       avl_tree_entry(_cur, struct_name,                 \
 323                                      struct_member), 1)                 \
 324                   && (_parent = avl_get_parent(_cur), 1);               \
 325              _cur = avl_tree_next_in_postorder(_cur, _parent))
 326
 327 #endif /* _AVL_TREE_H_ */