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