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