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  * Copyright 2022 Eric Biggers
   6  *
   7  * Permission is hereby granted, free of charge, to any person
   8  * obtaining a copy of this software and associated documentation
   9  * files (the "Software"), to deal in the Software without
  10  * restriction, including without limitation the rights to use,
  11  * copy, modify, merge, publish, distribute, sublicense, and/or sell
  12  * copies of the Software, and to permit persons to whom the
  13  * Software is furnished to do so, subject to the following
  14  * conditions:
  15  *
  16  * The above copyright notice and this permission notice shall be
  17  * included in all copies or substantial portions of the Software.
  18  *
  19  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  20  * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  21  * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  22  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  23  * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  24  * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  25  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  26  * OTHER DEALINGS IN THE SOFTWARE.
  27  */
  28
  29 #ifndef _AVL_TREE_H_
  30 #define _AVL_TREE_H_
  31
  32 #include "wimlib/types.h"
  33 #define AVL_INLINE forceinline
  34
  35 /* Node in an AVL tree.  Embed this in some other data structure.  */
  36 struct avl_tree_node {
  37
  38         /* Pointer to left child or NULL  */
  39         struct avl_tree_node *left;
  40
  41         /* Pointer to right child or NULL  */
  42         struct avl_tree_node *right;
  43
  44         /* Pointer to parent combined with the balance factor.  This saves 4 or
  45          * 8 bytes of memory depending on the CPU architecture.
  46          *
  47          * Low 2 bits:  One greater than the balance factor of this subtree,
  48          * which is equal to height(right) - height(left).  The mapping is:
  49          *
  50          * 00 => -1
  51          * 01 =>  0
  52          * 10 => +1
  53          * 11 => undefined
  54          *
  55          * The rest of the bits are the pointer to the parent node.  It must be
  56          * 4-byte aligned, and it will be NULL if this is the root node and
  57          * therefore has no parent.  */
  58         uintptr_t parent_balance;
  59 };
  60
  61 /* Cast an AVL tree node to the containing data structure.  */
  62 #define avl_tree_entry(entry, type, member) \
  63         ((type*) ((char *)(entry) - offsetof(type, member)))
  64
  65 /* Returns a pointer to the parent of the specified AVL tree node, or NULL if it
  66  * is already the root of the tree.  */
  67 static AVL_INLINE struct avl_tree_node *
  68 avl_get_parent(const struct avl_tree_node *node)
  69 {
  70         return (struct avl_tree_node *)(node->parent_balance & ~3);
  71 }
  72
  73 /* (Internal use only)  */
  74 extern void
  75 avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
  76                                 struct avl_tree_node *inserted);
  77
  78 /*
  79  * Looks up an item in the specified AVL tree.
  80  *
  81  * @root
  82  *      Pointer to the root of the AVL tree.  (This can be NULL --- that just
  83  *      means the tree is empty.)
  84  *
  85  * @cmp_ctx
  86  *      First argument to pass to the comparison callback.  This generally
  87  *      should be a pointer to an object equal to the one being searched for.
  88  *
  89  * @cmp
  90  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
  91  *      is less than, equal to, or greater than the second argument,
  92  *      respectively.  The first argument will be @cmp_ctx and the second
  93  *      argument will be a pointer to the AVL tree node of an item in the tree.
  94  *
  95  * Returns a pointer to the AVL tree node of the resulting item, or NULL if the
  96  * item was not found.
  97  *
  98  * Example:
  99  *
 100  * struct int_wrapper {
 101  *      int data;
 102  *      struct avl_tree_node index_node;
 103  * };
 104  *
 105  * static int _avl_cmp_int_to_node(const void *intptr,
 106  *                                 const struct avl_tree_node *nodeptr)
 107  * {
 108  *      int n1 = *(const int *)intptr;
 109  *      int n2 = avl_tree_entry(nodeptr, struct int_wrapper, index_node)->data;
 110  *      if (n1 < n2)
 111  *              return -1;
 112  *      else if (n1 > n2)
 113  *              return 1;
 114  *      else
 115  *              return 0;
 116  * }
 117  *
 118  * bool contains_int(struct avl_tree_node *root, int n)
 119  * {
 120  *      struct avl_tree_node *result;
 121  *
 122  *      result = avl_tree_lookup(root, &n, _avl_cmp_int_to_node);
 123  *      return result ? true : false;
 124  * }
 125  */
 126 static AVL_INLINE struct avl_tree_node *
 127 avl_tree_lookup(const struct avl_tree_node *root,
 128                 const void *cmp_ctx,
 129                 int (*cmp)(const void *, const struct avl_tree_node *))
 130 {
 131         const struct avl_tree_node *cur = root;
 132
 133         while (cur) {
 134                 int res = (*cmp)(cmp_ctx, cur);
 135                 if (res < 0)
 136                         cur = cur->left;
 137                 else if (res > 0)
 138                         cur = cur->right;
 139                 else
 140                         break;
 141         }
 142         return (struct avl_tree_node*)cur;
 143 }
 144
 145 /* Same as avl_tree_lookup(), but uses a more specific type for the comparison
 146  * function.  Specifically, with this function the item being searched for is
 147  * expected to be in the same format as those already in the tree, with an
 148  * embedded 'struct avl_tree_node'.  */
 149 static AVL_INLINE struct avl_tree_node *
 150 avl_tree_lookup_node(const struct avl_tree_node *root,
 151                      const struct avl_tree_node *node,
 152                      int (*cmp)(const struct avl_tree_node *,
 153                                 const struct avl_tree_node *))
 154 {
 155         const struct avl_tree_node *cur = root;
 156
 157         while (cur) {
 158                 int res = (*cmp)(node, cur);
 159                 if (res < 0)
 160                         cur = cur->left;
 161                 else if (res > 0)
 162                         cur = cur->right;
 163                 else
 164                         break;
 165         }
 166         return (struct avl_tree_node*)cur;
 167 }
 168
 169 /*
 170  * Inserts an item into the specified AVL tree.
 171  *
 172  * @root_ptr
 173  *      Location of the AVL tree's root pointer.  Indirection is needed because
 174  *      the root node may change as a result of rotations caused by the
 175  *      insertion.  Initialize *root_ptr to NULL for an empty tree.
 176  *
 177  * @item
 178  *      Pointer to the `struct avl_tree_node' embedded in the item to insert.
 179  *      No members in it need be pre-initialized, although members in the
 180  *      containing structure should be pre-initialized so that @cmp can use them
 181  *      in comparisons.
 182  *
 183  * @cmp
 184  *      Comparison callback.  Must return < 0, 0, or > 0 if the first argument
 185  *      is less than, equal to, or greater than the second argument,
 186  *      respectively.  The first argument will be @item and the second
 187  *      argument will be a pointer to an AVL tree node embedded in some
 188  *      previously-inserted item to which @item is being compared.
 189  *
 190  * If no item in the tree is comparatively equal (via @cmp) to @item, inserts
 191  * @item and returns NULL.  Otherwise does nothing and returns a pointer to the
 192  * AVL tree node embedded in the previously-inserted item which compared equal
 193  * to @item.
 194  *
 195  * Example:
 196  *
 197  * struct int_wrapper {
 198  *      int data;
 199  *      struct avl_tree_node index_node;
 200  * };
 201  *
 202  * #define GET_DATA(i) avl_tree_entry((i), struct int_wrapper, index_node)->data
 203  *
 204  * static int _avl_cmp_ints(const struct avl_tree_node *node1,
 205  *                          const struct avl_tree_node *node2)
 206  * {
 207  *      int n1 = GET_DATA(node1);
 208  *      int n2 = GET_DATA(node2);
 209  *      if (n1 < n2)
 210  *              return -1;
 211  *      else if (n1 > n2)
 212  *              return 1;
 213  *      else
 214  *              return 0;
 215  * }
 216  *
 217  * bool insert_int(struct avl_tree_node **root_ptr, int data)
 218  * {
 219  *      struct int_wrapper *i = malloc(sizeof(struct int_wrapper));
 220  *      i->data = data;
 221  *      if (avl_tree_insert(root_ptr, &i->index_node, _avl_cmp_ints)) {
 222  *              // Duplicate.
 223  *              free(i);
 224  *              return false;
 225  *      }
 226  *      return true;
 227  * }
 228  */
 229 static AVL_INLINE struct avl_tree_node *
 230 avl_tree_insert(struct avl_tree_node **root_ptr,
 231                 struct avl_tree_node *item,
 232                 int (*cmp)(const struct avl_tree_node *,
 233                            const struct avl_tree_node *))
 234 {
 235         struct avl_tree_node **cur_ptr = root_ptr, *cur = NULL;
 236         int res;
 237
 238         while (*cur_ptr) {
 239                 cur = *cur_ptr;
 240                 res = (*cmp)(item, cur);
 241                 if (res < 0)
 242                         cur_ptr = &cur->left;
 243                 else if (res > 0)
 244                         cur_ptr = &cur->right;
 245                 else
 246                         return cur;
 247         }
 248         *cur_ptr = item;
 249         item->parent_balance = (uintptr_t)cur | 1;
 250         avl_tree_rebalance_after_insert(root_ptr, item);
 251         return NULL;
 252 }
 253
 254 /* Removes an item from the specified AVL tree.
 255  * See implementation for details.  */
 256 extern void
 257 avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node);
 258
 259 /* Nonrecursive AVL tree traversal functions  */
 260
 261 extern struct avl_tree_node *
 262 avl_tree_first_in_order(const struct avl_tree_node *root);
 263
 264 extern struct avl_tree_node *
 265 avl_tree_last_in_order(const struct avl_tree_node *root);
 266
 267 extern struct avl_tree_node *
 268 avl_tree_next_in_order(const struct avl_tree_node *node);
 269
 270 extern struct avl_tree_node *
 271 avl_tree_prev_in_order(const struct avl_tree_node *node);
 272
 273 extern struct avl_tree_node *
 274 avl_tree_first_in_postorder(const struct avl_tree_node *root);
 275
 276 extern struct avl_tree_node *
 277 avl_tree_next_in_postorder(const struct avl_tree_node *prev,
 278                            const struct avl_tree_node *prev_parent);
 279
 280 /*
 281  * Iterate through the nodes in an AVL tree in sorted order.
 282  * You may not modify the tree during the iteration.
 283  *
 284  * @child_struct
 285  *      Variable that will receive a pointer to each struct inserted into the
 286  *      tree.
 287  * @root
 288  *      Root of the AVL tree.
 289  * @struct_name
 290  *      Type of *child_struct.
 291  * @struct_member
 292  *      Member of @struct_name type that is the AVL tree node.
 293  *
 294  * Example:
 295  *
 296  * struct int_wrapper {
 297  *      int data;
 298  *      struct avl_tree_node index_node;
 299  * };
 300  *
 301  * void print_ints(struct avl_tree_node *root)
 302  * {
 303  *      struct int_wrapper *i;
 304  *
 305  *      avl_tree_for_each_in_order(i, root, struct int_wrapper, index_node)
 306  *              printf("%d\n", i->data);
 307  * }
 308  */
 309 #define avl_tree_for_each_in_order(child_struct, root,                  \
 310                                    struct_name, struct_member)          \
 311         for (struct avl_tree_node *_cur =                               \
 312                 avl_tree_first_in_order(root);                          \
 313              _cur && ((child_struct) =                                  \
 314                       avl_tree_entry(_cur, struct_name,                 \
 315                                      struct_member), 1);                \
 316              _cur = avl_tree_next_in_order(_cur))
 317
 318 /*
 319  * Like avl_tree_for_each_in_order(), but uses the reverse order.
 320  */
 321 #define avl_tree_for_each_in_reverse_order(child_struct, root,          \
 322                                            struct_name, struct_member)  \
 323         for (struct avl_tree_node *_cur =                               \
 324                 avl_tree_last_in_order(root);                           \
 325              _cur && ((child_struct) =                                  \
 326                       avl_tree_entry(_cur, struct_name,                 \
 327                                      struct_member), 1);                \
 328              _cur = avl_tree_prev_in_order(_cur))
 329
 330 /*
 331  * Like avl_tree_for_each_in_order(), but iterates through the nodes in
 332  * postorder, so the current node may be deleted or freed.
 333  */
 334 #define avl_tree_for_each_in_postorder(child_struct, root,              \
 335                                        struct_name, struct_member)      \
 336         for (struct avl_tree_node *_cur =                               \
 337                 avl_tree_first_in_postorder(root), *_parent;            \
 338              _cur && ((child_struct) =                                  \
 339                       avl_tree_entry(_cur, struct_name,                 \
 340                                      struct_member), 1)                 \
 341                   && (_parent = avl_get_parent(_cur), 1);               \
 342              _cur = avl_tree_next_in_postorder(_cur, _parent))
 343
 344 #endif /* _AVL_TREE_H_ */