1 /*
2 * Copyright © 2017 Faith 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 set 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 static inline void
rb_tree_init(struct rb_tree * T)76 rb_tree_init(struct rb_tree *T)
77 {
78 T->root = NULL;
79 }
80
81
82 /** Returns true if the red-black tree is empty */
83 static inline bool
rb_tree_is_empty(const struct rb_tree * T)84 rb_tree_is_empty(const struct rb_tree *T)
85 {
86 return T->root == NULL;
87 }
88
89 /** Get the first (left-most) node in the tree or NULL */
90 struct rb_node *rb_tree_first(struct rb_tree *T);
91
92 /** Get the last (right-most) node in the tree or NULL */
93 struct rb_node *rb_tree_last(struct rb_tree *T);
94
95 /** Get the next node (to the right) in the tree or NULL */
96 struct rb_node *rb_node_next(struct rb_node *node);
97
98 /** Get the next previous (to the left) in the tree or NULL */
99 struct rb_node *rb_node_prev(struct rb_node *node);
100
101 #ifdef __cplusplus
102 /* This macro will not work correctly if `t' uses virtual inheritance. */
103 #define rb_tree_offsetof(t, f, p) \
104 (((char *) &((t *) p)->f) - ((char *) p))
105 #else
106 #define rb_tree_offsetof(t, f, p) offsetof(t, f)
107 #endif
108
109 /** Retrieve the data structure containing a node
110 *
111 * \param type The type of the containing data structure
112 *
113 * \param node A pointer to a rb_node
114 *
115 * \param field The rb_node field in the containing data structure
116 */
117 #define rb_node_data(type, node, field) \
118 ((type *)(((char *)(node)) - rb_tree_offsetof(type, field, node)))
119
120 /** Insert a node into a possibly augmented tree at a particular location
121 *
122 * This function should probably not be used directly as it relies on the
123 * caller to ensure that the parent node is correct. Use rb_tree_insert
124 * instead.
125 *
126 * If \p update is non-NULL, it will be called for the node being inserted as
127 * well as any nodes which have their children changed and all of their
128 * ancestors. The intent is that each node may contain some augmented data
129 * which is calculated recursively from the node itself and its children, and
130 * \p update should recalculate that data. It's assumed that the function used
131 * to calculate the node data is associative in order to avoid calling it
132 * redundantly after rebalancing the tree.
133 *
134 * \param T The red-black tree into which to insert the new node
135 *
136 * \param parent The node in the tree that will be the parent of the
137 * newly inserted node
138 *
139 * \param node The node to insert
140 *
141 * \param insert_left If true, the new node will be the left child of
142 * \p parent, otherwise it will be the right child
143 *
144 * \param update The optional function used to calculate per-node data
145 */
146 void rb_augmented_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
147 struct rb_node *node, bool insert_left,
148 void (*update)(struct rb_node *));
149
150 /** Insert a node into a tree at a particular location
151 *
152 * This function should probably not be used directly as it relies on the
153 * caller to ensure that the parent node is correct. Use rb_tree_insert
154 * instead.
155 *
156 * \param T The red-black tree into which to insert the new node
157 *
158 * \param parent The node in the tree that will be the parent of the
159 * newly inserted node
160 *
161 * \param node The node to insert
162 *
163 * \param insert_left If true, the new node will be the left child of
164 * \p parent, otherwise it will be the right child
165 */
166 static inline void
rb_tree_insert_at(struct rb_tree * T,struct rb_node * parent,struct rb_node * node,bool insert_left)167 rb_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
168 struct rb_node *node, bool insert_left)
169 {
170 rb_augmented_tree_insert_at(T, parent, node, insert_left, NULL);
171 }
172
173 /** Insert a node into a possibly augmented tree
174 *
175 * \param T The red-black tree into which to insert the new node
176 *
177 * \param node The node to insert
178 *
179 * \param cmp A comparison function to use to order the nodes.
180 *
181 * \param update Same meaning as in rb_augmented_tree_insert_at()
182 */
183 static inline void
rb_augmented_tree_insert(struct rb_tree * T,struct rb_node * node,int (* cmp)(const struct rb_node *,const struct rb_node *),void (* update)(struct rb_node *))184 rb_augmented_tree_insert(struct rb_tree *T, struct rb_node *node,
185 int (*cmp)(const struct rb_node *, const struct rb_node *),
186 void (*update)(struct rb_node *))
187 {
188 /* This function is declared inline in the hopes that the compiler can
189 * optimize away the comparison function pointer call.
190 */
191 struct rb_node *y = NULL;
192 struct rb_node *x = T->root;
193 bool left = false;
194 while (x != NULL) {
195 y = x;
196 left = cmp(x, node) < 0;
197 if (left)
198 x = x->left;
199 else
200 x = x->right;
201 }
202
203 rb_augmented_tree_insert_at(T, y, node, left, update);
204 }
205
206 /** Insert a node into a tree
207 *
208 * \param T The red-black tree into which to insert the new node
209 *
210 * \param node The node to insert
211 *
212 * \param cmp A comparison function to use to order the nodes.
213 */
214 static inline void
rb_tree_insert(struct rb_tree * T,struct rb_node * node,int (* cmp)(const struct rb_node *,const struct rb_node *))215 rb_tree_insert(struct rb_tree *T, struct rb_node *node,
216 int (*cmp)(const struct rb_node *, const struct rb_node *))
217 {
218 rb_augmented_tree_insert(T, node, cmp, NULL);
219 }
220
221 /** Remove a node from a possibly augmented tree
222 *
223 * \param T The red-black tree from which to remove the node
224 *
225 * \param node The node to remove
226 *
227 * \param update Same meaning as in rb_agumented_tree_insert_at()
228 *
229 */
230 void rb_augmented_tree_remove(struct rb_tree *T, struct rb_node *z,
231 void (*update)(struct rb_node *));
232
233 /** Remove a node from a tree
234 *
235 * \param T The red-black tree from which to remove the node
236 *
237 * \param node The node to remove
238 */
239 static inline void
rb_tree_remove(struct rb_tree * T,struct rb_node * z)240 rb_tree_remove(struct rb_tree *T, struct rb_node *z)
241 {
242 rb_augmented_tree_remove(T, z, NULL);
243 }
244
245 /** Search the tree for a node
246 *
247 * If a node with a matching key exists, the first matching node found will
248 * be returned. If no matching node exists, NULL is returned.
249 *
250 * \param T The red-black tree to search
251 *
252 * \param key The key to search for
253 *
254 * \param cmp A comparison function to use to order the nodes
255 */
256 static inline struct rb_node *
rb_tree_search(struct rb_tree * T,const void * key,int (* cmp)(const struct rb_node *,const void *))257 rb_tree_search(struct rb_tree *T, const void *key,
258 int (*cmp)(const struct rb_node *, const void *))
259 {
260 /* This function is declared inline in the hopes that the compiler can
261 * optimize away the comparison function pointer call.
262 */
263 struct rb_node *x = T->root;
264 while (x != NULL) {
265 int c = cmp(x, key);
266 if (c < 0) {
267 x = x->left;
268 } else if (c > 0) {
269 x = x->right;
270 } else {
271 /* x is the first *encountered* node matching the key. There may
272 * be other nodes in the left subtree that also match the key.
273 */
274 while (true) {
275 struct rb_node *prev = rb_node_prev(x);
276
277 if (prev == NULL || cmp(prev, key) != 0)
278 return x;
279
280 x = prev;
281 }
282 }
283 }
284
285 return x;
286 }
287
288 /** Sloppily search the tree for a node
289 *
290 * This function searches the tree for a given node. If a node with a
291 * matching key exists, that first encountered matching node found (there may
292 * be other matching nodes in the left subtree) will be returned. If no node
293 * with an exactly matching key exists, the node returned will be either the
294 * right-most node comparing less than \p key or the right-most node comparing
295 * greater than \p key. If the tree is empty, NULL is returned.
296 *
297 * \param T The red-black tree to search
298 *
299 * \param key The key to search for
300 *
301 * \param cmp A comparison function to use to order the nodes
302 */
303 static inline struct rb_node *
rb_tree_search_sloppy(struct rb_tree * T,const void * key,int (* cmp)(const struct rb_node *,const void *))304 rb_tree_search_sloppy(struct rb_tree *T, const void *key,
305 int (*cmp)(const struct rb_node *, const void *))
306 {
307 /* This function is declared inline in the hopes that the compiler can
308 * optimize away the comparison function pointer call.
309 */
310 struct rb_node *y = NULL;
311 struct rb_node *x = T->root;
312 while (x != NULL) {
313 y = x;
314 int c = cmp(x, key);
315 if (c < 0)
316 x = x->left;
317 else if (c > 0)
318 x = x->right;
319 else
320 return x;
321 }
322
323 return y;
324 }
325
326 #define rb_node_next_or_null(n) ((n) == NULL ? NULL : rb_node_next(n))
327 #define rb_node_prev_or_null(n) ((n) == NULL ? NULL : rb_node_prev(n))
328
329 /** Iterate over the nodes in the tree
330 *
331 * \param type The type of the containing data structure
332 *
333 * \param node The variable name for current node in the iteration;
334 * this will be declared as a pointer to \p type
335 *
336 * \param T The red-black tree
337 *
338 * \param field The rb_node field in containing data structure
339 */
340 #define rb_tree_foreach(type, iter, T, field) \
341 for (type *iter, *__node = (type *)rb_tree_first(T); \
342 __node != NULL && \
343 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
344 __node = (type *)rb_node_next((struct rb_node *)__node))
345
346 /** Iterate over the nodes in the tree, allowing the current node to be freed
347 *
348 * \param type The type of the containing data structure
349 *
350 * \param node The variable name for current node in the iteration;
351 * this will be declared as a pointer to \p type
352 *
353 * \param T The red-black tree
354 *
355 * \param field The rb_node field in containing data structure
356 */
357 #define rb_tree_foreach_safe(type, iter, T, field) \
358 for (type *iter, \
359 *__node = (type *)rb_tree_first(T), \
360 *__next = (type *)rb_node_next_or_null((struct rb_node *)__node); \
361 __node != NULL && \
362 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
363 __node = __next, \
364 __next = (type *)rb_node_next_or_null((struct rb_node *)__node))
365
366 /** Iterate over the nodes in the tree in reverse
367 *
368 * \param type The type of the containing data structure
369 *
370 * \param node The variable name for current node in the iteration;
371 * this will be declared as a pointer to \p type
372 *
373 * \param T The red-black tree
374 *
375 * \param field The rb_node field in containing data structure
376 */
377 #define rb_tree_foreach_rev(type, iter, T, field) \
378 for (type *iter, *__node = (type *)rb_tree_last(T); \
379 __node != NULL && \
380 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
381 __node = (type *)rb_node_prev((struct rb_node *)__node))
382
383 /** Iterate over the nodes in the tree in reverse, allowing the current node to be freed
384 *
385 * \param type The type of the containing data structure
386 *
387 * \param node The variable name for current node in the iteration;
388 * this will be declared as a pointer to \p type
389 *
390 * \param T The red-black tree
391 *
392 * \param field The rb_node field in containing data structure
393 */
394 #define rb_tree_foreach_rev_safe(type, iter, T, field) \
395 for (type *iter, \
396 *__node = (type *)rb_tree_last(T), \
397 *__prev = (type *)rb_node_prev_or_null((struct rb_node *)__node); \
398 __node != NULL && \
399 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
400 __node = __prev, \
401 __prev = (type *)rb_node_prev_or_null((struct rb_node *)__node))
402
403 /** Unsigned interval
404 *
405 * Intervals are closed by convention.
406 */
407 struct uinterval {
408 unsigned start, end;
409 };
410
411 struct uinterval_node {
412 struct rb_node node;
413
414 /* Must be filled in before inserting */
415 struct uinterval interval;
416
417 /* Managed internally by the tree */
418 unsigned max_end;
419 };
420
421 /** Insert a node into an unsigned interval tree. */
422 void uinterval_tree_insert(struct rb_tree *tree, struct uinterval_node *node);
423
424 /** Remove a node from an unsigned interval tree. */
425 void uinterval_tree_remove(struct rb_tree *tree, struct uinterval_node *node);
426
427 /** Get the first node intersecting the given interval. */
428 struct uinterval_node *uinterval_tree_first(struct rb_tree *tree,
429 struct uinterval interval);
430
431 /** Get the next node after \p node intersecting the given interval. */
432 struct uinterval_node *uinterval_node_next(struct uinterval_node *node,
433 struct uinterval interval);
434
435 /** Iterate over the nodes in the tree intersecting the given interval
436 *
437 * The iteration itself should take O(k log n) time, where k is the number of
438 * iterations of the loop and n is the size of the tree.
439 *
440 * \param type The type of the containing data structure
441 *
442 * \param node The variable name for current node in the iteration;
443 * this will be declared as a pointer to \p type
444 *
445 * \param interval The interval to be tested against.
446 *
447 * \param T The red-black tree
448 *
449 * \param field The uinterval_node field in containing data structure
450 */
451 #define uinterval_tree_foreach(type, iter, interval, T, field) \
452 for (type *iter, *__node = (type *)uinterval_tree_first(T, interval); \
453 __node != NULL && \
454 (iter = rb_node_data(type, (struct uinterval_node *)__node, field), true); \
455 __node = (type *)uinterval_node_next((struct uinterval_node *)__node, interval))
456
457 /** Validate a red-black tree
458 *
459 * This function walks the tree and validates that this is a valid red-
460 * black tree. If anything is wrong, it will assert-fail.
461 */
462 void rb_tree_validate(struct rb_tree *T);
463
464 #ifdef __cplusplus
465 } /* extern C */
466 #endif
467
468 #endif /* RB_TREE_H */
469