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