1/*
2 * Copyright © 2017 Jason Ekstrand
3 *
4 * Permission is hereby granted, free of charge, to any person obtaining a
5 * copy of this software and associated documentation files (the "Software"),
6 * to deal in the Software without restriction, including without limitation
7 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
8 * and/or sell copies of the Software, and to permit persons to whom the
9 * Software is furnished to do so, subject to the following conditions:
10 *
11 * The above copyright notice and this permission notice shall be included in
12 * all copies or substantial portions of the Software.
13 *
14 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
19 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
20 * DEALINGS IN THE SOFTWARE.
21 */
22
23#ifndef RB_TREE_H
24#define RB_TREE_H
25
26#include <stdbool.h>
27#include <stddef.h>
28#include <stdint.h>
29#include <stdlib.h>
30
31/** A red-black tree node
32 *
33 * This struct represents a node in the red-black tree.  This struct should
34 * be embedded as a member in whatever structure you wish to put in the
35 * tree.
36 */
37struct rb_node {
38    /** Parent and color of this node
39     *
40     * The least significant bit represents the color and is est to 1 for
41     * black and 0 for red.  The other bits are the pointer to the parent
42     * and that pointer can be retrieved by masking off the bottom bit and
43     * casting to a pointer.
44     */
45    uintptr_t parent;
46
47    /** Left child of this node, null for a leaf */
48    struct rb_node *left;
49
50    /** Right child of this node, null for a leaf */
51    struct rb_node *right;
52};
53
54/** Return the parent node of the given node or NULL if it is the root */
55static inline struct rb_node *
56rb_node_parent(struct rb_node *n)
57{
58    return (struct rb_node *)(n->parent & ~(uintptr_t)1);
59}
60
61/** A red-black tree
62 *
63 * This struct represents the red-black tree itself.  It is just a pointer
64 * to the root node with no other metadata.
65 */
66struct rb_tree {
67    struct rb_node *root;
68};
69
70/** Initialize a red-black tree */
71void rb_tree_init(struct rb_tree *T);
72
73/** Returns true if the red-black tree is empty */
74static inline bool
75rb_tree_is_empty(const struct rb_tree *T)
76{
77    return T->root == NULL;
78}
79
80/** Retrieve the data structure containing a node
81 *
82 * \param   type    The type of the containing data structure
83 *
84 * \param   node    A pointer to a rb_node
85 *
86 * \param   field   The rb_node field in the containing data structure
87 */
88#define rb_node_data(type, node, field) \
89    ((type *)(((char *)(node)) - offsetof(type, field)))
90
91/** Insert a node into a tree at a particular location
92 *
93 * This function should probably not be used directly as it relies on the
94 * caller to ensure that the parent node is correct.  Use rb_tree_insert
95 * instead.
96 *
97 * \param   T           The red-black tree into which to insert the new node
98 *
99 * \param   parent      The node in the tree that will be the parent of the
100 *                      newly inserted node
101 *
102 * \param   node        The node to insert
103 *
104 * \param   insert_left If true, the new node will be the left child of
105 *                      \p parent, otherwise it will be the right child
106 */
107void rb_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
108                       struct rb_node *node, bool insert_left);
109
110/** Insert a node into a tree
111 *
112 * \param   T       The red-black tree into which to insert the new node
113 *
114 * \param   node    The node to insert
115 *
116 * \param   cmp     A comparison function to use to order the nodes.
117 */
118static inline void
119rb_tree_insert(struct rb_tree *T, struct rb_node *node,
120               int (*cmp)(const struct rb_node *, const struct rb_node *))
121{
122    /* This function is declared inline in the hopes that the compiler can
123     * optimize away the comparison function pointer call.
124     */
125    struct rb_node *y = NULL;
126    struct rb_node *x = T->root;
127    bool left = false;
128    while (x != NULL) {
129        y = x;
130        left = cmp(node, x) < 0;
131        if (left)
132            x = x->left;
133        else
134            x = x->right;
135    }
136
137    rb_tree_insert_at(T, y, node, left);
138}
139
140/** Remove a node from a tree
141 *
142 * \param   T       The red-black tree from which to remove the node
143 *
144 * \param   node    The node to remove
145 */
146void rb_tree_remove(struct rb_tree *T, struct rb_node *z);
147
148/** Search the tree for a node
149 *
150 * If a node with a matching key exists, the first matching node found will
151 * be returned.  If no matching node exists, NULL is returned.
152 *
153 * \param   T       The red-black tree to search
154 *
155 * \param   key     The key to search for
156 *
157 * \param   cmp     A comparison function to use to order the nodes
158 */
159static inline struct rb_node *
160rb_tree_search(struct rb_tree *T, const void *key,
161               int (*cmp)(const struct rb_node *, const void *))
162{
163    /* This function is declared inline in the hopes that the compiler can
164     * optimize away the comparison function pointer call.
165     */
166    struct rb_node *x = T->root;
167    while (x != NULL) {
168        int c = cmp(x, key);
169        if (c < 0)
170            x = x->right;
171        else if (c > 0)
172            x = x->left;
173        else
174            return x;
175    }
176
177    return x;
178}
179
180/** Sloppily search the tree for a node
181 *
182 * This function searches the tree for a given node.  If a node with a
183 * matching key exists, that first matching node found will be returned.
184 * If no node with an exactly matching key exists, the node returned will
185 * be either the right-most node comparing less than \p key or the
186 * right-most node comparing greater than \p key.  If the tree is empty,
187 * NULL is returned.
188 *
189 * \param   T       The red-black tree to search
190 *
191 * \param   key     The key to search for
192 *
193 * \param   cmp     A comparison function to use to order the nodes
194 */
195static inline struct rb_node *
196rb_tree_search_sloppy(struct rb_tree *T, const void *key,
197                      int (*cmp)(const struct rb_node *, const void *))
198{
199    /* This function is declared inline in the hopes that the compiler can
200     * optimize away the comparison function pointer call.
201     */
202    struct rb_node *y = NULL;
203    struct rb_node *x = T->root;
204    while (x != NULL) {
205        y = x;
206        int c = cmp(x, key);
207        if (c < 0)
208            x = x->right;
209        else if (c > 0)
210            x = x->left;
211        else
212            return x;
213    }
214
215    return y;
216}
217
218/** Get the first (left-most) node in the tree or NULL */
219struct rb_node *rb_tree_first(struct rb_tree *T);
220
221/** Get the last (right-most) node in the tree or NULL */
222struct rb_node *rb_tree_last(struct rb_tree *T);
223
224/** Get the next node (to the right) in the tree or NULL */
225struct rb_node *rb_node_next(struct rb_node *node);
226
227/** Get the next previous (to the left) in the tree or NULL */
228struct rb_node *rb_node_prev(struct rb_node *node);
229
230/** Get the next node if available or the same node again.
231 *
232 * \param   type    The type of the containing data structure
233 *
234 * \param   node    The variable name for current node in the iteration;
235 *                  this will be declared as a pointer to \p type
236 *
237 * \param   field   The rb_node field in containing data structure
238 */
239#define rb_tree_node_next_if_available(type, node, field) \
240   (&node->field != NULL) ? rb_node_data(type, rb_node_next(&node->field), field) : node
241
242/** Get the previous node if available or the same node again.
243 *
244 * \param   type    The type of the containing data structure
245 *
246 * \param   node    The variable name for current node in the iteration;
247 *                  this will be declared as a pointer to \p type
248 *
249 * \param   field   The rb_node field in containing data structure
250 */
251#define rb_tree_node_prev_if_available(type, node, field) \
252   (&node->field != NULL) ? rb_node_data(type, rb_node_prev(&node->field), field) : node
253
254/** Iterate over the nodes in the tree
255 *
256 * \param   type    The type of the containing data structure
257 *
258 * \param   node    The variable name for current node in the iteration;
259 *                  this will be declared as a pointer to \p type
260 *
261 * \param   T       The red-black tree
262 *
263 * \param   field   The rb_node field in containing data structure
264 */
265#define rb_tree_foreach(type, node, T, field) \
266   for (type *node = rb_node_data(type, rb_tree_first(T), field); \
267        &node->field != NULL; \
268        node = rb_node_data(type, rb_node_next(&node->field), field))
269
270/** Iterate over the nodes in the tree, allowing the current node to be freed
271 *
272 * \param   type    The type of the containing data structure
273 *
274 * \param   node    The variable name for current node in the iteration;
275 *                  this will be declared as a pointer to \p type
276 *
277 * \param   T       The red-black tree
278 *
279 * \param   field   The rb_node field in containing data structure
280 */
281#define rb_tree_foreach_safe(type, node, T, field) \
282   for (type *node = rb_node_data(type, rb_tree_first(T), field), \
283           *__next = rb_tree_node_next_if_available(type, node, field); \
284        &node->field != NULL; \
285        node = __next, __next = rb_tree_node_next_if_available(type, node, field))
286
287/** Iterate over the nodes in the tree in reverse
288 *
289 * \param   type    The type of the containing data structure
290 *
291 * \param   node    The variable name for current node in the iteration;
292 *                  this will be declared as a pointer to \p type
293 *
294 * \param   T       The red-black tree
295 *
296 * \param   field   The rb_node field in containing data structure
297 */
298#define rb_tree_foreach_rev(type, node, T, field) \
299   for (type *node = rb_node_data(type, rb_tree_last(T), field); \
300        &node->field != NULL; \
301        node = rb_node_data(type, rb_node_prev(&node->field), field))
302
303/** Iterate over the nodes in the tree in reverse, allowing the current node to be freed
304 *
305 * \param   type    The type of the containing data structure
306 *
307 * \param   node    The variable name for current node in the iteration;
308 *                  this will be declared as a pointer to \p type
309 *
310 * \param   T       The red-black tree
311 *
312 * \param   field   The rb_node field in containing data structure
313 */
314#define rb_tree_foreach_rev_safe(type, node, T, field) \
315   for (type *node = rb_node_data(type, rb_tree_last(T), field), \
316           *__prev = rb_tree_node_prev_if_available(type, node, field);  \
317        &node->field != NULL; \
318        node = __prev, __prev = rb_tree_node_prev_if_available(type, node, field))
319
320/** Validate a red-black tree
321 *
322 * This function walks the tree and validates that this is a valid red-
323 * black tree.  If anything is wrong, it will assert-fail.
324 */
325void rb_tree_validate(struct rb_tree *T);
326
327#endif /* RB_TREE_H */
328