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         const struct avl_tree_node *cur = root;
 149
 150         while (cur) {
 151                 int res = (*cmp)(node, cur);
 152                 if (res < 0)
 153                         cur = cur->left;
 154                 else if (res > 0)
 155                         cur = cur->right;
 156                 else
 157                         break;
 158         }
 159         return (struct avl_tree_node*)cur;
 160 }
 161
 162 /*
 163  * Inserts an item into the specified AVL tree.
 164  *
 165  * @root_ptr
 166  *      Location of the AVL tree's root pointer.  Indirection is needed because
 167  *      the root node may change as a result of rotations caused by the
 168  *      insertion.  Initialize *root_ptr to NULL for an empty tree.
 169  *
 170  * @item
 171  *      Pointer to the `struct avl_tree_node' embedded in the item to insert.
 172  *      No members in it need be pre-initialized, although members in the
 173  *      containing structure should be pre-initialized so that @cmp can use them
 174  *      in comparisons.
 175  *
 176  * @cmp
 177  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
 178  *      is less than, equal to, or greater than the second argument,
 179  *      respectively.  The first argument will be @item and the second
 180  *      argument will be a pointer to an AVL tree node embedded in some
 181  *      previously-inserted item to which @item is being compared.
 182  *
 183  * If no item in the tree is comparatively equal (via @cmp) to @item, inserts
 184  * @item and returns NULL.  Otherwise does nothing and returns a pointer to the
 185  * AVL tree node embedded in the previously-inserted item which compared equal
 186  * to @item.
 187  *
 188  * Example:
 189  *
 190  * struct int_wrapper {
 191  *      int data;
 192  *      struct avl_tree_node index_node;
 193  * };
 194  *
 195  * #define GET_DATA(i) avl_tree_entry((i), struct int_wrapper, index_node)->data
 196  *
 197  * static int _avl_cmp_ints(const struct avl_tree_node *node1,
 198  *                          const struct avl_tree_node *node2)
 199  * {
 200  *      int n1 = GET_DATA(node1);
 201  *      int n2 = GET_DATA(node2);
 202  *      if (n1 < n2)
 203  *              return -1;
 204  *      else if (n1 > n2)
 205  *              return 1;
 206  *      else
 207  *              return 0;
 208  * }
 209  *
 210  * bool insert_int(struct avl_tree_node **root_ptr, int data)
 211  * {
 212  *      struct int_wrapper *i = malloc(sizeof(struct int_wrapper));
 213  *      i->data = data;
 214  *      if (avl_tree_insert(root_ptr, &i->index_node, _avl_cmp_ints)) {
 215  *              // Duplicate.
 216  *              free(i);
 217  *              return false;
 218  *      }
 219  *      return true;
 220  * }
 221  */
 222 static AVL_INLINE struct avl_tree_node *
 223 avl_tree_insert(struct avl_tree_node **root_ptr,
 224                 struct avl_tree_node *item,
 225                 int (*cmp)(const struct avl_tree_node *,
 226                            const struct avl_tree_node *))
 227 {
 228         struct avl_tree_node **cur_ptr = root_ptr, *cur = NULL;
 229         int res;
 230
 231         while (*cur_ptr) {
 232                 cur = *cur_ptr;
 233                 res = (*cmp)(item, cur);
 234                 if (res < 0)
 235                         cur_ptr = &cur->left;
 236                 else if (res > 0)
 237                         cur_ptr = &cur->right;
 238                 else
 239                         return cur;
 240         }
 241         *cur_ptr = item;
 242         item->parent_balance = (uintptr_t)cur | 1;
 243         avl_tree_rebalance_after_insert(root_ptr, item);
 244         return NULL;
 245 }
 246
 247 /* Removes an item from the specified AVL tree.
 248  * See implementation for details.  */
 249 extern void
 250 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node);
 251
 252 /* Nonrecursive AVL tree traversal functions  */
 253
 254 extern struct avl_tree_node *
 255 avl_tree_first_in_order(const struct avl_tree_node *root);
 256
 257 extern struct avl_tree_node *
 258 avl_tree_last_in_order(const struct avl_tree_node *root);
 259
 260 extern struct avl_tree_node *
 261 avl_tree_next_in_order(const struct avl_tree_node *node);
 262
 263 extern struct avl_tree_node *
 264 avl_tree_prev_in_order(const struct avl_tree_node *node);
 265
 266 extern struct avl_tree_node *
 267 avl_tree_first_in_postorder(const struct avl_tree_node *root);
 268
 269 extern struct avl_tree_node *
 270 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
 271                            const struct avl_tree_node *prev_parent);
 272
 273 /*
 274  * Iterate through the nodes in an AVL tree in sorted order.
 275  * You may not modify the tree during the iteration.
 276  *
 277  * @child_struct
 278  *      Variable that will receive a pointer to each struct inserted into the
 279  *      tree.
 280  * @root
 281  *      Root of the AVL tree.
 282  * @struct_name
 283  *      Type of *child_struct.
 284  * @struct_member
 285  *      Member of @struct_name type that is the AVL tree node.
 286  *
 287  * Example:
 288  *
 289  * struct int_wrapper {
 290  *      int data;
 291  *      struct avl_tree_node index_node;
 292  * };
 293  *
 294  * void print_ints(struct avl_tree_node *root)
 295  * {
 296  *      struct int_wrapper *i;
 297  *
 298  *      avl_tree_for_each_in_order(i, root, struct int_wrapper, index_node)
 299  *              printf("%d\n", i->data);
 300  * }
 301  */
 302 #define avl_tree_for_each_in_order(child_struct, root,                  \
 303                                    struct_name, struct_member)          \
 304         for (struct avl_tree_node *_cur =                               \
 305                 avl_tree_first_in_order(root);                          \
 306              _cur && ((child_struct) =                                  \
 307                       avl_tree_entry(_cur, struct_name,                 \
 308                                      struct_member), 1);                \
 309              _cur = avl_tree_next_in_order(_cur))
 310
 311 /*
 312  * Like avl_tree_for_each_in_order(), but uses the reverse order.
 313  */
 314 #define avl_tree_for_each_in_reverse_order(child_struct, root,          \
 315                                            struct_name, struct_member)  \
 316         for (struct avl_tree_node *_cur =                               \
 317                 avl_tree_last_in_order(root);                           \
 318              _cur && ((child_struct) =                                  \
 319                       avl_tree_entry(_cur, struct_name,                 \
 320                                      struct_member), 1);                \
 321              _cur = avl_tree_prev_in_order(_cur))
 322
 323 /*
 324  * Like avl_tree_for_each_in_order(), but iterates through the nodes in
 325  * postorder, so the current node may be deleted or freed.
 326  */
 327 #define avl_tree_for_each_in_postorder(child_struct, root,              \
 328                                        struct_name, struct_member)      \
 329         for (struct avl_tree_node *_cur =                               \
 330                 avl_tree_first_in_postorder(root), *_parent;            \
 331              _cur && ((child_struct) =                                  \
 332                       avl_tree_entry(_cur, struct_name,                 \
 333                                      struct_member), 1)                 \
 334                   && (_parent = avl_get_parent(_cur), 1);               \
 335              _cur = avl_tree_next_in_postorder(_cur, _parent))
 336
 337 #endif /* _AVL_TREE_H_ */