1 /*
2 * =============================================================================
3 *
4 * Filename: rbtree.c
5 *
6 * Description: rbtree(Red-Black tree) implementation adapted from linux
7 * kernel thus can be used in userspace c program.
8 *
9 * Created: 09/02/2012 11:38:12 PM
10 *
11 * Author: Fu Haiping (forhappy), haipingf@gmail.com
12 * Company: ICT ( Institute Of Computing Technology, CAS )
13 *
14 * =============================================================================
15 */
16
17 /*
18 Red Black Trees
19 (C) 1999 Andrea Arcangeli <andrea@suse.de>
20 (C) 2002 David Woodhouse <dwmw2@infradead.org>
21
22 This program is free software; you can redistribute it and/or modify
23 it under the terms of the GNU General Public License as published by
24 the Free Software Foundation; either version 2 of the License, or
25 (at your option) any later version.
26
27 This program is distributed in the hope that it will be useful,
28 but WITHOUT ANY WARRANTY; without even the implied warranty of
29 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
30 GNU General Public License for more details.
31
32 You should have received a copy of the GNU General Public License
33 along with this program; if not, write to the Free Software
34 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
35
36 linux/lib/rbtree.c
37 */
38
39 #include "general.h"
40 #include "rbtree.h"
41
__rb_rotate_left(struct rb_node * node,struct rb_root * root)42 static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
43 {
44 struct rb_node *right = node->rb_right;
45 struct rb_node *parent = rb_parent(node);
46
47 if ((node->rb_right = right->rb_left))
48 rb_set_parent(right->rb_left, node);
49 right->rb_left = node;
50
51 rb_set_parent(right, parent);
52
53 if (parent)
54 {
55 if (node == parent->rb_left)
56 parent->rb_left = right;
57 else
58 parent->rb_right = right;
59 }
60 else
61 root->rb_node = right;
62 rb_set_parent(node, right);
63 }
64
__rb_rotate_right(struct rb_node * node,struct rb_root * root)65 static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
66 {
67 struct rb_node *left = node->rb_left;
68 struct rb_node *parent = rb_parent(node);
69
70 if ((node->rb_left = left->rb_right))
71 rb_set_parent(left->rb_right, node);
72 left->rb_right = node;
73
74 rb_set_parent(left, parent);
75
76 if (parent)
77 {
78 if (node == parent->rb_right)
79 parent->rb_right = left;
80 else
81 parent->rb_left = left;
82 }
83 else
84 root->rb_node = left;
85 rb_set_parent(node, left);
86 }
87
rb_insert_color(struct rb_node * node,struct rb_root * root)88 void rb_insert_color(struct rb_node *node, struct rb_root *root)
89 {
90 struct rb_node *parent, *gparent;
91
92 while ((parent = rb_parent(node)) && rb_is_red(parent))
93 {
94 gparent = rb_parent(parent);
95
96 if (parent == gparent->rb_left)
97 {
98 {
99 register struct rb_node *uncle = gparent->rb_right;
100 if (uncle && rb_is_red(uncle))
101 {
102 rb_set_black(uncle);
103 rb_set_black(parent);
104 rb_set_red(gparent);
105 node = gparent;
106 continue;
107 }
108 }
109
110 if (parent->rb_right == node)
111 {
112 register struct rb_node *tmp;
113 __rb_rotate_left(parent, root);
114 tmp = parent;
115 parent = node;
116 node = tmp;
117 }
118
119 rb_set_black(parent);
120 rb_set_red(gparent);
121 __rb_rotate_right(gparent, root);
122 } else {
123 {
124 register struct rb_node *uncle = gparent->rb_left;
125 if (uncle && rb_is_red(uncle))
126 {
127 rb_set_black(uncle);
128 rb_set_black(parent);
129 rb_set_red(gparent);
130 node = gparent;
131 continue;
132 }
133 }
134
135 if (parent->rb_left == node)
136 {
137 register struct rb_node *tmp;
138 __rb_rotate_right(parent, root);
139 tmp = parent;
140 parent = node;
141 node = tmp;
142 }
143
144 rb_set_black(parent);
145 rb_set_red(gparent);
146 __rb_rotate_left(gparent, root);
147 }
148 }
149
150 rb_set_black(root->rb_node);
151 }
152
__rb_erase_color(struct rb_node * node,struct rb_node * parent,struct rb_root * root)153 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
154 struct rb_root *root)
155 {
156 struct rb_node *other;
157
158 while ((!node || rb_is_black(node)) && node != root->rb_node)
159 {
160 if (parent->rb_left == node)
161 {
162 other = parent->rb_right;
163 if (rb_is_red(other))
164 {
165 rb_set_black(other);
166 rb_set_red(parent);
167 __rb_rotate_left(parent, root);
168 other = parent->rb_right;
169 }
170 if ((!other->rb_left || rb_is_black(other->rb_left)) &&
171 (!other->rb_right || rb_is_black(other->rb_right)))
172 {
173 rb_set_red(other);
174 node = parent;
175 parent = rb_parent(node);
176 }
177 else
178 {
179 if (!other->rb_right || rb_is_black(other->rb_right))
180 {
181 rb_set_black(other->rb_left);
182 rb_set_red(other);
183 __rb_rotate_right(other, root);
184 other = parent->rb_right;
185 }
186 rb_set_color(other, rb_color(parent));
187 rb_set_black(parent);
188 rb_set_black(other->rb_right);
189 __rb_rotate_left(parent, root);
190 node = root->rb_node;
191 break;
192 }
193 }
194 else
195 {
196 other = parent->rb_left;
197 if (rb_is_red(other))
198 {
199 rb_set_black(other);
200 rb_set_red(parent);
201 __rb_rotate_right(parent, root);
202 other = parent->rb_left;
203 }
204 if ((!other->rb_left || rb_is_black(other->rb_left)) &&
205 (!other->rb_right || rb_is_black(other->rb_right)))
206 {
207 rb_set_red(other);
208 node = parent;
209 parent = rb_parent(node);
210 }
211 else
212 {
213 if (!other->rb_left || rb_is_black(other->rb_left))
214 {
215 rb_set_black(other->rb_right);
216 rb_set_red(other);
217 __rb_rotate_left(other, root);
218 other = parent->rb_left;
219 }
220 rb_set_color(other, rb_color(parent));
221 rb_set_black(parent);
222 rb_set_black(other->rb_left);
223 __rb_rotate_right(parent, root);
224 node = root->rb_node;
225 break;
226 }
227 }
228 }
229 if (node)
230 rb_set_black(node);
231 }
232
rb_erase(struct rb_node * node,struct rb_root * root)233 void rb_erase(struct rb_node *node, struct rb_root *root)
234 {
235 struct rb_node *child, *parent;
236 int color;
237
238 if (!node->rb_left)
239 child = node->rb_right;
240 else if (!node->rb_right)
241 child = node->rb_left;
242 else
243 {
244 struct rb_node *old = node, *left;
245
246 node = node->rb_right;
247 while ((left = node->rb_left) != NULL)
248 node = left;
249
250 if (rb_parent(old)) {
251 if (rb_parent(old)->rb_left == old)
252 rb_parent(old)->rb_left = node;
253 else
254 rb_parent(old)->rb_right = node;
255 } else
256 root->rb_node = node;
257
258 child = node->rb_right;
259 parent = rb_parent(node);
260 color = rb_color(node);
261
262 if (parent == old) {
263 parent = node;
264 } else {
265 if (child)
266 rb_set_parent(child, parent);
267 parent->rb_left = child;
268
269 node->rb_right = old->rb_right;
270 rb_set_parent(old->rb_right, node);
271 }
272
273 node->rb_parent_color = old->rb_parent_color;
274 node->rb_left = old->rb_left;
275 rb_set_parent(old->rb_left, node);
276
277 goto color;
278 }
279
280 parent = rb_parent(node);
281 color = rb_color(node);
282
283 if (child)
284 rb_set_parent(child, parent);
285 if (parent)
286 {
287 if (parent->rb_left == node)
288 parent->rb_left = child;
289 else
290 parent->rb_right = child;
291 }
292 else
293 root->rb_node = child;
294
295 color:
296 if (color == RB_BLACK)
297 __rb_erase_color(child, parent, root);
298 }
299
rb_augment_path(struct rb_node * node,rb_augment_f func,void * data)300 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
301 {
302 struct rb_node *parent;
303
304 up:
305 func(node, data);
306 parent = rb_parent(node);
307 if (!parent)
308 return;
309
310 if (node == parent->rb_left && parent->rb_right)
311 func(parent->rb_right, data);
312 else if (parent->rb_left)
313 func(parent->rb_left, data);
314
315 node = parent;
316 goto up;
317 }
318
319 /*
320 * after inserting @node into the tree, update the tree to account for
321 * both the new entry and any damage done by rebalance
322 */
rb_augment_insert(struct rb_node * node,rb_augment_f func,void * data)323 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
324 {
325 if (node->rb_left)
326 node = node->rb_left;
327 else if (node->rb_right)
328 node = node->rb_right;
329
330 rb_augment_path(node, func, data);
331 }
332
333 /*
334 * before removing the node, find the deepest node on the rebalance path
335 * that will still be there after @node gets removed
336 */
rb_augment_erase_begin(struct rb_node * node)337 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
338 {
339 struct rb_node *deepest;
340
341 if (!node->rb_right && !node->rb_left)
342 deepest = rb_parent(node);
343 else if (!node->rb_right)
344 deepest = node->rb_left;
345 else if (!node->rb_left)
346 deepest = node->rb_right;
347 else {
348 deepest = rb_next(node);
349 if (deepest->rb_right)
350 deepest = deepest->rb_right;
351 else if (rb_parent(deepest) != node)
352 deepest = rb_parent(deepest);
353 }
354
355 return deepest;
356 }
357
358 /*
359 * after removal, update the tree to account for the removed entry
360 * and any rebalance damage.
361 */
rb_augment_erase_end(struct rb_node * node,rb_augment_f func,void * data)362 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
363 {
364 if (node)
365 rb_augment_path(node, func, data);
366 }
367
368 /*
369 * This function returns the first node (in sort order) of the tree.
370 */
rb_first(const struct rb_root * root)371 struct rb_node *rb_first(const struct rb_root *root)
372 {
373 struct rb_node *n;
374
375 n = root->rb_node;
376 if (!n)
377 return NULL;
378 while (n->rb_left)
379 n = n->rb_left;
380 return n;
381 }
382
rb_last(const struct rb_root * root)383 struct rb_node *rb_last(const struct rb_root *root)
384 {
385 struct rb_node *n;
386
387 n = root->rb_node;
388 if (!n)
389 return NULL;
390 while (n->rb_right)
391 n = n->rb_right;
392 return n;
393 }
394
rb_next(const struct rb_node * node)395 struct rb_node *rb_next(const struct rb_node *node)
396 {
397 struct rb_node *parent;
398
399 if (rb_parent(node) == node)
400 return NULL;
401
402 /* If we have a right-hand child, go down and then left as far
403 as we can. */
404 if (node->rb_right) {
405 node = node->rb_right;
406 while (node->rb_left)
407 node=node->rb_left;
408 return (struct rb_node *)node;
409 }
410
411 /* No right-hand children. Everything down and left is
412 smaller than us, so any 'next' node must be in the general
413 direction of our parent. Go up the tree; any time the
414 ancestor is a right-hand child of its parent, keep going
415 up. First time it's a left-hand child of its parent, said
416 parent is our 'next' node. */
417 while ((parent = rb_parent(node)) && node == parent->rb_right)
418 node = parent;
419
420 return parent;
421 }
422
rb_prev(const struct rb_node * node)423 struct rb_node *rb_prev(const struct rb_node *node)
424 {
425 struct rb_node *parent;
426
427 if (rb_parent(node) == node)
428 return NULL;
429
430 /* If we have a left-hand child, go down and then right as far
431 as we can. */
432 if (node->rb_left) {
433 node = node->rb_left;
434 while (node->rb_right)
435 node=node->rb_right;
436 return (struct rb_node *)node;
437 }
438
439 /* No left-hand children. Go up till we find an ancestor which
440 is a right-hand child of its parent */
441 while ((parent = rb_parent(node)) && node == parent->rb_left)
442 node = parent;
443
444 return parent;
445 }
446
rb_replace_node(struct rb_node * victim,struct rb_node * new,struct rb_root * root)447 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
448 struct rb_root *root)
449 {
450 struct rb_node *parent = rb_parent(victim);
451
452 /* Set the surrounding nodes to point to the replacement */
453 if (parent) {
454 if (victim == parent->rb_left)
455 parent->rb_left = new;
456 else
457 parent->rb_right = new;
458 } else {
459 root->rb_node = new;
460 }
461 if (victim->rb_left)
462 rb_set_parent(victim->rb_left, new);
463 if (victim->rb_right)
464 rb_set_parent(victim->rb_right, new);
465
466 /* Copy the pointers/colour from the victim to the replacement */
467 *new = *victim;
468 }
469