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