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

   1 /*
   2  * avl_tree.h - intrusive, nonrecursive AVL tree data structure (self-balancing
   3  *              binary search tree), header file
   4  *
   5  * The following copying information applies to this specific source code file:
   6  *
   7  * Written in 2014-2016 by Eric Biggers <ebiggers3@gmail.com>
   8  *
   9  * To the extent possible under law, the author(s) have dedicated all copyright
  10  * and related and neighboring rights to this software to the public domain
  11  * worldwide via the Creative Commons Zero 1.0 Universal Public Domain
  12  * Dedication (the "CC0").
  13  *
  14  * This software is distributed in the hope that it will be useful, but WITHOUT
  15  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  16  * FOR A PARTICULAR PURPOSE. See the CC0 for more details.
  17  *
  18  * You should have received a copy of the CC0 along with this software; if not
  19  * see <http://creativecommons.org/publicdomain/zero/1.0/>.
  20  */
  21
  22 #ifndef _AVL_TREE_H_
  23 #define _AVL_TREE_H_
  24
  25 #include "wimlib/types.h"
  26 #define AVL_INLINE forceinline
  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 /* (Internal use only)  */
  67 extern void
  68 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
  69                                 struct avl_tree_node *inserted);
  70
  71 /*
  72  * Looks up an item in the specified AVL tree.
  73  *
  74  * @root
  75  *      Pointer to the root of the AVL tree.  (This can be NULL --- that just
  76  *      means the tree is empty.)
  77  *
  78  * @cmp_ctx
  79  *      First argument to pass to the comparison callback.  This generally
  80  *      should be a pointer to an object equal to the one being searched for.
  81  *
  82  * @cmp
  83  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
  84  *      is less than, equal to, or greater than the second argument,
  85  *      respectively.  The first argument will be @cmp_ctx and the second
  86  *      argument will be a pointer to the AVL tree node of an item in the tree.
  87  *
  88  * Returns a pointer to the AVL tree node of the resulting item, or NULL if the
  89  * item was not found.
  90  *
  91  * Example:
  92  *
  93  * struct int_wrapper {
  94  *      int data;
  95  *      struct avl_tree_node index_node;
  96  * };
  97  *
  98  * static int _avl_cmp_int_to_node(const void *intptr,
  99  *                                 const struct avl_tree_node *nodeptr)
 100  * {
 101  *      int n1 = *(const int *)intptr;
 102  *      int n2 = avl_tree_entry(nodeptr, struct int_wrapper, index_node)->data;
 103  *      if (n1 < n2)
 104  *              return -1;
 105  *      else if (n1 > n2)
 106  *              return 1;
 107  *      else
 108  *              return 0;
 109  * }
 110  *
 111  * bool contains_int(struct avl_tree_node *root, int n)
 112  * {
 113  *      struct avl_tree_node *result;
 114  *
 115  *      result = avl_tree_lookup(root, &n, _avl_cmp_int_to_node);
 116  *      return result ? true : false;
 117  * }
 118  */
 119 static AVL_INLINE struct avl_tree_node *
 120 avl_tree_lookup(const struct avl_tree_node *root,
 121                 const void *cmp_ctx,
 122                 int (*cmp)(const void *, const struct avl_tree_node *))
 123 {
 124         const struct avl_tree_node *cur = root;
 125
 126         while (cur) {
 127                 int res = (*cmp)(cmp_ctx, cur);
 128                 if (res < 0)
 129                         cur = cur->left;
 130                 else if (res > 0)
 131                         cur = cur->right;
 132                 else
 133                         break;
 134         }
 135         return (struct avl_tree_node*)cur;
 136 }
 137
 138 /* Same as avl_tree_lookup(), but uses a more specific type for the comparison
 139  * function.  Specifically, with this function the item being searched for is
 140  * expected to be in the same format as those already in the tree, with an
 141  * embedded 'struct avl_tree_node'.  */
 142 static AVL_INLINE struct avl_tree_node *
 143 avl_tree_lookup_node(const struct avl_tree_node *root,
 144                      const struct avl_tree_node *node,
 145                      int (*cmp)(const struct avl_tree_node *,
 146                                 const struct avl_tree_node *))
 147 {
 148         return avl_tree_lookup(root,
 149                                (const void *)node,
 150                                (int (*) (const void *,
 151                                          const struct avl_tree_node *))cmp);
 152 }
 153
 154 /*
 155  * Inserts an item into the specified AVL tree.
 156  *
 157  * @root_ptr
 158  *      Location of the AVL tree's root pointer.  Indirection is needed because
 159  *      the root node may change as a result of rotations caused by the
 160  *      insertion.  Initialize *root_ptr to NULL for an empty tree.
 161  *
 162  * @item
 163  *      Pointer to the `struct avl_tree_node' embedded in the item to insert.
 164  *      No members in it need be pre-initialized, although members in the
 165  *      containing structure should be pre-initialized so that @cmp can use them
 166  *      in comparisons.
 167  *
 168  * @cmp
 169  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
 170  *      is less than, equal to, or greater than the second argument,
 171  *      respectively.  The first argument will be @item and the second
 172  *      argument will be a pointer to an AVL tree node embedded in some
 173  *      previously-inserted item to which @item is being compared.
 174  *
 175  * If no item in the tree is comparatively equal (via @cmp) to @item, inserts
 176  * @item and returns NULL.  Otherwise does nothing and returns a pointer to the
 177  * AVL tree node embedded in the previously-inserted item which compared equal
 178  * to @item.
 179  *
 180  * Example:
 181  *
 182  * struct int_wrapper {
 183  *      int data;
 184  *      struct avl_tree_node index_node;
 185  * };
 186  *
 187  * #define GET_DATA(i) avl_tree_entry((i), struct int_wrapper, index_node)->data
 188  *
 189  * static int _avl_cmp_ints(const struct avl_tree_node *node1,
 190  *                          const struct avl_tree_node *node2)
 191  * {
 192  *      int n1 = GET_DATA(node1);
 193  *      int n2 = GET_DATA(node2);
 194  *      if (n1 < n2)
 195  *              return -1;
 196  *      else if (n1 > n2)
 197  *              return 1;
 198  *      else
 199  *              return 0;
 200  * }
 201  *
 202  * bool insert_int(struct avl_tree_node **root_ptr, int data)
 203  * {
 204  *      struct int_wrapper *i = malloc(sizeof(struct int_wrapper));
 205  *      i->data = data;
 206  *      if (avl_tree_insert(root_ptr, &i->index_node, _avl_cmp_ints)) {
 207  *              // Duplicate.
 208  *              free(i);
 209  *              return false;
 210  *      }
 211  *      return true;
 212  * }
 213  */
 214 static AVL_INLINE struct avl_tree_node *
 215 avl_tree_insert(struct avl_tree_node **root_ptr,
 216                 struct avl_tree_node *item,
 217                 int (*cmp)(const struct avl_tree_node *,
 218                            const struct avl_tree_node *))
 219 {
 220         struct avl_tree_node **cur_ptr = root_ptr, *cur = NULL;
 221         int res;
 222
 223         while (*cur_ptr) {
 224                 cur = *cur_ptr;
 225                 res = (*cmp)(item, cur);
 226                 if (res < 0)
 227                         cur_ptr = &cur->left;
 228                 else if (res > 0)
 229                         cur_ptr = &cur->right;
 230                 else
 231                         return cur;
 232         }
 233         *cur_ptr = item;
 234         item->parent_balance = (uintptr_t)cur | 1;
 235         avl_tree_rebalance_after_insert(root_ptr, item);
 236         return NULL;
 237 }
 238
 239 /* Removes an item from the specified AVL tree.
 240  * See implementation for details.  */
 241 extern void
 242 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node);
 243
 244 /* Nonrecursive AVL tree traversal functions  */
 245
 246 extern struct avl_tree_node *
 247 avl_tree_first_in_order(const struct avl_tree_node *root);
 248
 249 extern struct avl_tree_node *
 250 avl_tree_last_in_order(const struct avl_tree_node *root);
 251
 252 extern struct avl_tree_node *
 253 avl_tree_next_in_order(const struct avl_tree_node *node);
 254
 255 extern struct avl_tree_node *
 256 avl_tree_prev_in_order(const struct avl_tree_node *node);
 257
 258 extern struct avl_tree_node *
 259 avl_tree_first_in_postorder(const struct avl_tree_node *root);
 260
 261 extern struct avl_tree_node *
 262 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
 263                            const struct avl_tree_node *prev_parent);
 264
 265 /*
 266  * Iterate through the nodes in an AVL tree in sorted order.
 267  * You may not modify the tree during the iteration.
 268  *
 269  * @child_struct
 270  *      Variable that will receive a pointer to each struct inserted into the
 271  *      tree.
 272  * @root
 273  *      Root of the AVL tree.
 274  * @struct_name
 275  *      Type of *child_struct.
 276  * @struct_member
 277  *      Member of @struct_name type that is the AVL tree node.
 278  *
 279  * Example:
 280  *
 281  * struct int_wrapper {
 282  *      int data;
 283  *      struct avl_tree_node index_node;
 284  * };
 285  *
 286  * void print_ints(struct avl_tree_node *root)
 287  * {
 288  *      struct int_wrapper *i;
 289  *
 290  *      avl_tree_for_each_in_order(i, root, struct int_wrapper, index_node)
 291  *              printf("%d\n", i->data);
 292  * }
 293  */
 294 #define avl_tree_for_each_in_order(child_struct, root,                  \
 295                                    struct_name, struct_member)          \
 296         for (struct avl_tree_node *_cur =                               \
 297                 avl_tree_first_in_order(root);                          \
 298              _cur && ((child_struct) =                                  \
 299                       avl_tree_entry(_cur, struct_name,                 \
 300                                      struct_member), 1);                \
 301              _cur = avl_tree_next_in_order(_cur))
 302
 303 /*
 304  * Like avl_tree_for_each_in_order(), but uses the reverse order.
 305  */
 306 #define avl_tree_for_each_in_reverse_order(child_struct, root,          \
 307                                            struct_name, struct_member)  \
 308         for (struct avl_tree_node *_cur =                               \
 309                 avl_tree_last_in_order(root);                           \
 310              _cur && ((child_struct) =                                  \
 311                       avl_tree_entry(_cur, struct_name,                 \
 312                                      struct_member), 1);                \
 313              _cur = avl_tree_prev_in_order(_cur))
 314
 315 /*
 316  * Like avl_tree_for_each_in_order(), but iterates through the nodes in
 317  * postorder, so the current node may be deleted or freed.
 318  */
 319 #define avl_tree_for_each_in_postorder(child_struct, root,              \
 320                                        struct_name, struct_member)      \
 321         for (struct avl_tree_node *_cur =                               \
 322                 avl_tree_first_in_postorder(root), *_parent;            \
 323              _cur && ((child_struct) =                                  \
 324                       avl_tree_entry(_cur, struct_name,                 \
 325                                      struct_member), 1)                 \
 326                   && (_parent = avl_get_parent(_cur), 1);               \
 327              _cur = avl_tree_next_in_postorder(_cur, _parent))
 328
 329 #endif /* _AVL_TREE_H_ */