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 #include "rb_tree.h"
24
25 /** \file rb_tree.c
26 *
27 * An implementation of a red-black tree
28 *
29 * This file implements the guts of a red-black tree. The implementation
30 * is mostly based on the one in "Introduction to Algorithms", third
31 * edition, by Cormen, Leiserson, Rivest, and Stein. The primary
32 * divergence in our algorithms from those presented in CLRS is that we use
33 * NULL for the leaves instead of a sentinel. This means we have to do a
34 * tiny bit more tracking in our implementation of delete but it makes the
35 * algorithms far more explicit than stashing stuff in the sentinel.
36 */
37
38 #include <stdlib.h>
39 #include <string.h>
40 #include <assert.h>
41
42 static bool
rb_node_is_black(struct rb_node * n)43 rb_node_is_black(struct rb_node *n)
44 {
45 /* NULL nodes are leaves and therefore black */
46 return (n == NULL) || (n->parent & 1);
47 }
48
49 static bool
rb_node_is_red(struct rb_node * n)50 rb_node_is_red(struct rb_node *n)
51 {
52 return !rb_node_is_black(n);
53 }
54
55 static void
rb_node_set_black(struct rb_node * n)56 rb_node_set_black(struct rb_node *n)
57 {
58 n->parent |= 1;
59 }
60
61 static void
rb_node_set_red(struct rb_node * n)62 rb_node_set_red(struct rb_node *n)
63 {
64 n->parent &= ~1ull;
65 }
66
67 static void
rb_node_copy_color(struct rb_node * dst,struct rb_node * src)68 rb_node_copy_color(struct rb_node *dst, struct rb_node *src)
69 {
70 dst->parent = (dst->parent & ~1ull) | (src->parent & 1);
71 }
72
73 static void
rb_node_set_parent(struct rb_node * n,struct rb_node * p)74 rb_node_set_parent(struct rb_node *n, struct rb_node *p)
75 {
76 n->parent = (n->parent & 1) | (uintptr_t)p;
77 }
78
79 static struct rb_node *
rb_node_minimum(struct rb_node * node)80 rb_node_minimum(struct rb_node *node)
81 {
82 while (node->left)
83 node = node->left;
84 return node;
85 }
86
87 static struct rb_node *
rb_node_maximum(struct rb_node * node)88 rb_node_maximum(struct rb_node *node)
89 {
90 while (node->right)
91 node = node->right;
92 return node;
93 }
94
95 void
rb_tree_init(struct rb_tree * T)96 rb_tree_init(struct rb_tree *T)
97 {
98 T->root = NULL;
99 }
100
101 /**
102 * Replace the subtree of T rooted at u with the subtree rooted at v
103 *
104 * This is called RB-transplant in CLRS.
105 *
106 * The node to be replaced is assumed to be a non-leaf.
107 */
108 static void
rb_tree_splice(struct rb_tree * T,struct rb_node * u,struct rb_node * v)109 rb_tree_splice(struct rb_tree *T, struct rb_node *u, struct rb_node *v)
110 {
111 assert(u);
112 struct rb_node *p = rb_node_parent(u);
113 if (p == NULL) {
114 assert(T->root == u);
115 T->root = v;
116 } else if (u == p->left) {
117 p->left = v;
118 } else {
119 assert(u == p->right);
120 p->right = v;
121 }
122 if (v)
123 rb_node_set_parent(v, p);
124 }
125
126 static void
rb_tree_rotate_left(struct rb_tree * T,struct rb_node * x)127 rb_tree_rotate_left(struct rb_tree *T, struct rb_node *x)
128 {
129 assert(x && x->right);
130
131 struct rb_node *y = x->right;
132 x->right = y->left;
133 if (y->left)
134 rb_node_set_parent(y->left, x);
135 rb_tree_splice(T, x, y);
136 y->left = x;
137 rb_node_set_parent(x, y);
138 }
139
140 static void
rb_tree_rotate_right(struct rb_tree * T,struct rb_node * y)141 rb_tree_rotate_right(struct rb_tree *T, struct rb_node *y)
142 {
143 assert(y && y->left);
144
145 struct rb_node *x = y->left;
146 y->left = x->right;
147 if (x->right)
148 rb_node_set_parent(x->right, y);
149 rb_tree_splice(T, y, x);
150 x->right = y;
151 rb_node_set_parent(y, x);
152 }
153
154 void
rb_tree_insert_at(struct rb_tree * T,struct rb_node * parent,struct rb_node * node,bool insert_left)155 rb_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
156 struct rb_node *node, bool insert_left)
157 {
158 /* This sets null children, parent, and a color of red */
159 memset(node, 0, sizeof(*node));
160
161 if (parent == NULL) {
162 assert(T->root == NULL);
163 T->root = node;
164 rb_node_set_black(node);
165 return;
166 }
167
168 if (insert_left) {
169 assert(parent->left == NULL);
170 parent->left = node;
171 } else {
172 assert(parent->right == NULL);
173 parent->right = node;
174 }
175 rb_node_set_parent(node, parent);
176
177 /* Now we do the insertion fixup */
178 struct rb_node *z = node;
179 while (rb_node_is_red(rb_node_parent(z))) {
180 struct rb_node *z_p = rb_node_parent(z);
181 assert(z == z_p->left || z == z_p->right);
182 struct rb_node *z_p_p = rb_node_parent(z_p);
183 assert(z_p_p != NULL);
184 if (z_p == z_p_p->left) {
185 struct rb_node *y = z_p_p->right;
186 if (rb_node_is_red(y)) {
187 rb_node_set_black(z_p);
188 rb_node_set_black(y);
189 rb_node_set_red(z_p_p);
190 z = z_p_p;
191 } else {
192 if (z == z_p->right) {
193 z = z_p;
194 rb_tree_rotate_left(T, z);
195 /* We changed z */
196 z_p = rb_node_parent(z);
197 assert(z == z_p->left || z == z_p->right);
198 z_p_p = rb_node_parent(z_p);
199 }
200 rb_node_set_black(z_p);
201 rb_node_set_red(z_p_p);
202 rb_tree_rotate_right(T, z_p_p);
203 }
204 } else {
205 struct rb_node *y = z_p_p->left;
206 if (rb_node_is_red(y)) {
207 rb_node_set_black(z_p);
208 rb_node_set_black(y);
209 rb_node_set_red(z_p_p);
210 z = z_p_p;
211 } else {
212 if (z == z_p->left) {
213 z = z_p;
214 rb_tree_rotate_right(T, z);
215 /* We changed z */
216 z_p = rb_node_parent(z);
217 assert(z == z_p->left || z == z_p->right);
218 z_p_p = rb_node_parent(z_p);
219 }
220 rb_node_set_black(z_p);
221 rb_node_set_red(z_p_p);
222 rb_tree_rotate_left(T, z_p_p);
223 }
224 }
225 }
226 rb_node_set_black(T->root);
227 }
228
229 void
rb_tree_remove(struct rb_tree * T,struct rb_node * z)230 rb_tree_remove(struct rb_tree *T, struct rb_node *z)
231 {
232 /* x_p is always the parent node of X. We have to track this
233 * separately because x may be NULL.
234 */
235 struct rb_node *x, *x_p;
236 struct rb_node *y = z;
237 bool y_was_black = rb_node_is_black(y);
238 if (z->left == NULL) {
239 x = z->right;
240 x_p = rb_node_parent(z);
241 rb_tree_splice(T, z, x);
242 } else if (z->right == NULL) {
243 x = z->left;
244 x_p = rb_node_parent(z);
245 rb_tree_splice(T, z, x);
246 } else {
247 /* Find the minimum sub-node of z->right */
248 y = rb_node_minimum(z->right);
249 y_was_black = rb_node_is_black(y);
250
251 x = y->right;
252 if (rb_node_parent(y) == z) {
253 x_p = y;
254 } else {
255 x_p = rb_node_parent(y);
256 rb_tree_splice(T, y, x);
257 y->right = z->right;
258 rb_node_set_parent(y->right, y);
259 }
260 assert(y->left == NULL);
261 rb_tree_splice(T, z, y);
262 y->left = z->left;
263 rb_node_set_parent(y->left, y);
264 rb_node_copy_color(y, z);
265 }
266
267 assert(x_p == NULL || x == x_p->left || x == x_p->right);
268
269 if (!y_was_black)
270 return;
271
272 /* Fixup RB tree after the delete */
273 while (x != T->root && rb_node_is_black(x)) {
274 if (x == x_p->left) {
275 struct rb_node *w = x_p->right;
276 if (rb_node_is_red(w)) {
277 rb_node_set_black(w);
278 rb_node_set_red(x_p);
279 rb_tree_rotate_left(T, x_p);
280 assert(x == x_p->left);
281 w = x_p->right;
282 }
283 if (rb_node_is_black(w->left) && rb_node_is_black(w->right)) {
284 rb_node_set_red(w);
285 x = x_p;
286 } else {
287 if (rb_node_is_black(w->right)) {
288 rb_node_set_black(w->left);
289 rb_node_set_red(w);
290 rb_tree_rotate_right(T, w);
291 w = x_p->right;
292 }
293 rb_node_copy_color(w, x_p);
294 rb_node_set_black(x_p);
295 rb_node_set_black(w->right);
296 rb_tree_rotate_left(T, x_p);
297 x = T->root;
298 }
299 } else {
300 struct rb_node *w = x_p->left;
301 if (rb_node_is_red(w)) {
302 rb_node_set_black(w);
303 rb_node_set_red(x_p);
304 rb_tree_rotate_right(T, x_p);
305 assert(x == x_p->right);
306 w = x_p->left;
307 }
308 if (rb_node_is_black(w->right) && rb_node_is_black(w->left)) {
309 rb_node_set_red(w);
310 x = x_p;
311 } else {
312 if (rb_node_is_black(w->left)) {
313 rb_node_set_black(w->right);
314 rb_node_set_red(w);
315 rb_tree_rotate_left(T, w);
316 w = x_p->left;
317 }
318 rb_node_copy_color(w, x_p);
319 rb_node_set_black(x_p);
320 rb_node_set_black(w->left);
321 rb_tree_rotate_right(T, x_p);
322 x = T->root;
323 }
324 }
325 x_p = rb_node_parent(x);
326 }
327 if (x)
328 rb_node_set_black(x);
329 }
330
331 struct rb_node *
rb_tree_first(struct rb_tree * T)332 rb_tree_first(struct rb_tree *T)
333 {
334 return T->root ? rb_node_minimum(T->root) : NULL;
335 }
336
337 struct rb_node *
rb_tree_last(struct rb_tree * T)338 rb_tree_last(struct rb_tree *T)
339 {
340 return T->root ? rb_node_maximum(T->root) : NULL;
341 }
342
343 struct rb_node *
rb_node_next(struct rb_node * node)344 rb_node_next(struct rb_node *node)
345 {
346 if (node->right) {
347 /* If we have a right child, then the next thing (compared to this
348 * node) is the left-most child of our right child.
349 */
350 return rb_node_minimum(node->right);
351 } else {
352 /* If node doesn't have a right child, crawl back up the to the
353 * left until we hit a parent to the right.
354 */
355 struct rb_node *p = rb_node_parent(node);
356 while (p && node == p->right) {
357 node = p;
358 p = rb_node_parent(node);
359 }
360 assert(p == NULL || node == p->left);
361 return p;
362 }
363 }
364
365 struct rb_node *
rb_node_prev(struct rb_node * node)366 rb_node_prev(struct rb_node *node)
367 {
368 if (node->left) {
369 /* If we have a left child, then the previous thing (compared to
370 * this node) is the right-most child of our left child.
371 */
372 return rb_node_maximum(node->left);
373 } else {
374 /* If node doesn't have a left child, crawl back up the to the
375 * right until we hit a parent to the left.
376 */
377 struct rb_node *p = rb_node_parent(node);
378 while (p && node == p->left) {
379 node = p;
380 p = rb_node_parent(node);
381 }
382 assert(p == NULL || node == p->right);
383 return p;
384 }
385 }
386
387 static void
validate_rb_node(struct rb_node * n,int black_depth)388 validate_rb_node(struct rb_node *n, int black_depth)
389 {
390 if (n == NULL) {
391 assert(black_depth == 0);
392 return;
393 }
394
395 if (rb_node_is_black(n)) {
396 black_depth--;
397 } else {
398 assert(rb_node_is_black(n->left));
399 assert(rb_node_is_black(n->right));
400 }
401
402 validate_rb_node(n->left, black_depth);
403 validate_rb_node(n->right, black_depth);
404 }
405
406 void
rb_tree_validate(struct rb_tree * T)407 rb_tree_validate(struct rb_tree *T)
408 {
409 if (T->root == NULL)
410 return;
411
412 assert(rb_node_is_black(T->root));
413
414 unsigned black_depth = 0;
415 for (struct rb_node *n = T->root; n; n = n->left) {
416 if (rb_node_is_black(n))
417 black_depth++;
418 }
419
420 validate_rb_node(T->root, black_depth);
421 }
422