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 */
37 struct 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 */
55 static inline struct rb_node *
rb_node_parent(struct rb_node * n)56 rb_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 */
66 struct rb_tree {
67 struct rb_node *root;
68 };
69
70 /** Initialize a red-black tree */
71 void rb_tree_init(struct rb_tree *T);
72
73 /** Returns true if the red-black tree is empty */
74 static inline bool
rb_tree_is_empty(const struct rb_tree * T)75 rb_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 */
107 void 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 */
118 static inline void
rb_tree_insert(struct rb_tree * T,struct rb_node * node,int (* cmp)(const struct rb_node *,const struct rb_node *))119 rb_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(x, node) < 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 */
146 void 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 */
159 static inline struct rb_node *
rb_tree_search(struct rb_tree * T,const void * key,int (* cmp)(const struct rb_node *,const void *))160 rb_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->left;
171 else if (c > 0)
172 x = x->right;
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 */
195 static inline struct rb_node *
rb_tree_search_sloppy(struct rb_tree * T,const void * key,int (* cmp)(const struct rb_node *,const void *))196 rb_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->left;
209 else if (c > 0)
210 x = x->right;
211 else
212 return x;
213 }
214
215 return y;
216 }
217
218 /** Get the first (left-most) node in the tree or NULL */
219 struct rb_node *rb_tree_first(struct rb_tree *T);
220
221 /** Get the last (right-most) node in the tree or NULL */
222 struct rb_node *rb_tree_last(struct rb_tree *T);
223
224 /** Get the next node (to the right) in the tree or NULL */
225 struct rb_node *rb_node_next(struct rb_node *node);
226
227 /** Get the next previous (to the left) in the tree or NULL */
228 struct rb_node *rb_node_prev(struct rb_node *node);
229
230 #define rb_node_next_or_null(n) ((n) == NULL ? NULL : rb_node_next(n))
231 #define rb_node_prev_or_null(n) ((n) == NULL ? NULL : rb_node_prev(n))
232
233 /** Iterate over the nodes in the tree
234 *
235 * \param type The type of the containing data structure
236 *
237 * \param node The variable name for current node in the iteration;
238 * this will be declared as a pointer to \p type
239 *
240 * \param T The red-black tree
241 *
242 * \param field The rb_node field in containing data structure
243 */
244 #define rb_tree_foreach(type, iter, T, field) \
245 for (type *iter, *__node = (type *)rb_tree_first(T); \
246 __node != NULL && \
247 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
248 __node = (type *)rb_node_next((struct rb_node *)__node))
249
250 /** Iterate over the nodes in the tree, allowing the current node to be freed
251 *
252 * \param type The type of the containing data structure
253 *
254 * \param node The variable name for current node in the iteration;
255 * this will be declared as a pointer to \p type
256 *
257 * \param T The red-black tree
258 *
259 * \param field The rb_node field in containing data structure
260 */
261 #define rb_tree_foreach_safe(type, iter, T, field) \
262 for (type *iter, \
263 *__node = (type *)rb_tree_first(T), \
264 *__next = (type *)rb_node_next_or_null((struct rb_node *)__node); \
265 __node != NULL && \
266 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
267 __node = __next, \
268 __next = (type *)rb_node_next_or_null((struct rb_node *)__node))
269
270 /** Iterate over the nodes in the tree in reverse
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_rev(type, iter, T, field) \
282 for (type *iter, *__node = (type *)rb_tree_last(T); \
283 __node != NULL && \
284 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
285 __node = (type *)rb_node_prev((struct rb_node *)__node))
286
287 /** Iterate over the nodes in the tree in reverse, allowing the current node to be freed
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_safe(type, iter, T, field) \
299 for (type *iter, \
300 *__node = (type *)rb_tree_last(T), \
301 *__prev = (type *)rb_node_prev_or_null((struct rb_node *)__node); \
302 __node != NULL && \
303 (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
304 __node = __prev, \
305 __prev = (type *)rb_node_prev_or_null((struct rb_node *)__node))
306
307 /** Validate a red-black tree
308 *
309 * This function walks the tree and validates that this is a valid red-
310 * black tree. If anything is wrong, it will assert-fail.
311 */
312 void rb_tree_validate(struct rb_tree *T);
313
314 #endif /* RB_TREE_H */
315