• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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