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 #ifdef __cplusplus
32 extern "C" {
33 #endif
34
35 /** A red-black tree node
36 *
37 * This struct represents a node in the red-black tree. This struct should
38 * be embedded as a member in whatever structure you wish to put in the
39 * tree.
40 */
41 struct rb_node {
42 /** Parent and color of this node
43 *
44 * The least significant bit represents the color and is est to 1 for
45 * black and 0 for red. The other bits are the pointer to the parent
46 * and that pointer can be retrieved by masking off the bottom bit and
47 * casting to a pointer.
48 */
49 uintptr_t parent;
50
51 /** Left child of this node, null for a leaf */
52 struct rb_node *left;
53
54 /** Right child of this node, null for a leaf */
55 struct rb_node *right;
56 };
57
58 /** Return the parent node of the given node or NULL if it is the root */
59 static inline struct rb_node *
rb_node_parent(struct rb_node * n)60 rb_node_parent(struct rb_node *n)
61 {
62 return (struct rb_node *)(n->parent & ~(uintptr_t)1);
63 }
64
65 /** A red-black tree
66 *
67 * This struct represents the red-black tree itself. It is just a pointer
68 * to the root node with no other metadata.
69 */
70 struct rb_tree {
71 struct rb_node *root;
72 };
73
74 /** Initialize a red-black tree */
75 void rb_tree_init(struct rb_tree *T);
76
77 /** Returns true if the red-black tree is empty */
78 static inline bool
rb_tree_is_empty(const struct rb_tree * T)79 rb_tree_is_empty(const struct rb_tree *T)
80 {
81 return T->root == NULL;
82 }
83
84 /** Retrieve the data structure containing a node
85 *
86 * \param type The type of the containing data structure
87 *
88 * \param node A pointer to a rb_node
89 *
90 * \param field The rb_node field in the containing data structure
91 */
92 #define rb_node_data(type, node, field) \
93 ((type *)(((char *)(node)) - offsetof(type, field)))
94
95 /** Insert a node into a tree at a particular location
96 *
97 * This function should probably not be used directly as it relies on the
98 * caller to ensure that the parent node is correct. Use rb_tree_insert
99 * instead.
100 *
101 * \param T The red-black tree into which to insert the new node
102 *
103 * \param parent The node in the tree that will be the parent of the
104 * newly inserted node
105 *
106 * \param node The node to insert
107 *
108 * \param insert_left If true, the new node will be the left child of
109 * \p parent, otherwise it will be the right child
110 */
111 void rb_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
112 struct rb_node *node, bool insert_left);
113
114 /** Insert a node into a tree
115 *
116 * \param T The red-black tree into which to insert the new node
117 *
118 * \param node The node to insert
119 *
120 * \param cmp A comparison function to use to order the nodes.
121 */
122 static inline void
rb_tree_insert(struct rb_tree * T,struct rb_node * node,int (* cmp)(const struct rb_node *,const struct rb_node *))123 rb_tree_insert(struct rb_tree *T, struct rb_node *node,
124 int (*cmp)(const struct rb_node *, const struct rb_node *))
125 {
126 /* This function is declared inline in the hopes that the compiler can
127 * optimize away the comparison function pointer call.
128 */
129 struct rb_node *y = NULL;
130 struct rb_node *x = T->root;
131 bool left = false;
132 while (x != NULL) {
133 y = x;
134 left = cmp(x, node) < 0;
135 if (left)
136 x = x->left;
137 else
138 x = x->right;
139 }
140
141 rb_tree_insert_at(T, y, node, left);
142 }
143
144 /** Remove a node from a tree
145 *
146 * \param T The red-black tree from which to remove the node
147 *
148 * \param node The node to remove
149 */
150 void rb_tree_remove(struct rb_tree *T, struct rb_node *z);
151
152 /** Search the tree for a node
153 *
154 * If a node with a matching key exists, the first matching node found will
155 * be returned. If no matching node exists, NULL is returned.
156 *
157 * \param T The red-black tree to search
158 *
159 * \param key The key to search for
160 *
161 * \param cmp A comparison function to use to order the nodes
162 */
163 static inline struct rb_node *
rb_tree_search(struct rb_tree * T,const void * key,int (* cmp)(const struct rb_node *,const void *))164 rb_tree_search(struct rb_tree *T, const void *key,
165 int (*cmp)(const struct rb_node *, const void *))
166 {
167 /* This function is declared inline in the hopes that the compiler can
168 * optimize away the comparison function pointer call.
169 */
170 struct rb_node *x = T->root;
171 while (x != NULL) {
172 int c = cmp(x, key);
173 if (c < 0)
174 x = x->left;
175 else if (c > 0)
176 x = x->right;
177 else
178 return x;
179 }
180
181 return x;
182 }
183
184 /** Sloppily search the tree for a node
185 *
186 * This function searches the tree for a given node. If a node with a
187 * matching key exists, that first matching node found will be returned.
188 * If no node with an exactly matching key exists, the node returned will
189 * be either the right-most node comparing less than \p key or the
190 * right-most node comparing greater than \p key. If the tree is empty,
191 * NULL is returned.
192 *
193 * \param T The red-black tree to search
194 *
195 * \param key The key to search for
196 *
197 * \param cmp A comparison function to use to order the nodes
198 */
199 static inline struct rb_node *
rb_tree_search_sloppy(struct rb_tree * T,const void * key,int (* cmp)(const struct rb_node *,const void *))200 rb_tree_search_sloppy(struct rb_tree *T, const void *key,
201 int (*cmp)(const struct rb_node *, const void *))
202 {
203 /* This function is declared inline in the hopes that the compiler can
204 * optimize away the comparison function pointer call.
205 */
206 struct rb_node *y = NULL;
207 struct rb_node *x = T->root;
208 while (x != NULL) {
209 y = x;
210 int c = cmp(x, key);
211 if (c < 0)
212 x = x->left;
213 else if (c > 0)
214 x = x->right;
215 else
216 return x;
217 }
218
219 return y;
220 }
221
222 /** Get the first (left-most) node in the tree or NULL */
223 struct rb_node *rb_tree_first(struct rb_tree *T);
224
225 /** Get the last (right-most) node in the tree or NULL */
226 struct rb_node *rb_tree_last(struct rb_tree *T);
227
228 /** Get the next node (to the right) in the tree or NULL */
229 struct rb_node *rb_node_next(struct rb_node *node);
230
231 /** Get the next previous (to the left) in the tree or NULL */
232 struct rb_node *rb_node_prev(struct rb_node *node);
233
234 #define rb_node_next_or_null(n) ((n) == NULL ? NULL : rb_node_next(n))
235 #define rb_node_prev_or_null(n) ((n) == NULL ? NULL : rb_node_prev(n))
236
237 /** Iterate over the nodes in the tree
238 *
239 * \param type The type of the containing data structure
240 *
241 * \param node The variable name for current node in the iteration;
242 * this will be declared as a pointer to \p type
243 *
244 * \param T The red-black tree
245 *
246 * \param field The rb_node field in containing data structure
247 */
248 #define rb_tree_foreach(type, iter, T, field) \
249 for (type *iter, *__node = (type *)rb_tree_first(T); \
250 __node != NULL && \
251 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
252 __node = (type *)rb_node_next((struct rb_node *)__node))
253
254 /** Iterate over the nodes in the tree, allowing the current node to be freed
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_safe(type, iter, T, field) \
266 for (type *iter, \
267 *__node = (type *)rb_tree_first(T), \
268 *__next = (type *)rb_node_next_or_null((struct rb_node *)__node); \
269 __node != NULL && \
270 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
271 __node = __next, \
272 __next = (type *)rb_node_next_or_null((struct rb_node *)__node))
273
274 /** Iterate over the nodes in the tree in reverse
275 *
276 * \param type The type of the containing data structure
277 *
278 * \param node The variable name for current node in the iteration;
279 * this will be declared as a pointer to \p type
280 *
281 * \param T The red-black tree
282 *
283 * \param field The rb_node field in containing data structure
284 */
285 #define rb_tree_foreach_rev(type, iter, T, field) \
286 for (type *iter, *__node = (type *)rb_tree_last(T); \
287 __node != NULL && \
288 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
289 __node = (type *)rb_node_prev((struct rb_node *)__node))
290
291 /** Iterate over the nodes in the tree in reverse, allowing the current node to be freed
292 *
293 * \param type The type of the containing data structure
294 *
295 * \param node The variable name for current node in the iteration;
296 * this will be declared as a pointer to \p type
297 *
298 * \param T The red-black tree
299 *
300 * \param field The rb_node field in containing data structure
301 */
302 #define rb_tree_foreach_rev_safe(type, iter, T, field) \
303 for (type *iter, \
304 *__node = (type *)rb_tree_last(T), \
305 *__prev = (type *)rb_node_prev_or_null((struct rb_node *)__node); \
306 __node != NULL && \
307 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
308 __node = __prev, \
309 __prev = (type *)rb_node_prev_or_null((struct rb_node *)__node))
310
311 /** Validate a red-black tree
312 *
313 * This function walks the tree and validates that this is a valid red-
314 * black tree. If anything is wrong, it will assert-fail.
315 */
316 void rb_tree_validate(struct rb_tree *T);
317
318 #ifdef __cplusplus
319 } /* extern C */
320 #endif
321
322 #endif /* RB_TREE_H */
323