xref: /linux/lib/rbtree.c (revision f8324e20f8289dffc646d64366332e05eaacab25)
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/module.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 	if (root->augment_cb) {
49 		root->augment_cb(node);
50 		root->augment_cb(right);
51 	}
52 }
53 
54 static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
55 {
56 	struct rb_node *left = node->rb_left;
57 	struct rb_node *parent = rb_parent(node);
58 
59 	if ((node->rb_left = left->rb_right))
60 		rb_set_parent(left->rb_right, node);
61 	left->rb_right = node;
62 
63 	rb_set_parent(left, parent);
64 
65 	if (parent)
66 	{
67 		if (node == parent->rb_right)
68 			parent->rb_right = left;
69 		else
70 			parent->rb_left = left;
71 	}
72 	else
73 		root->rb_node = left;
74 	rb_set_parent(node, left);
75 
76 	if (root->augment_cb) {
77 		root->augment_cb(node);
78 		root->augment_cb(left);
79 	}
80 }
81 
82 void rb_insert_color(struct rb_node *node, struct rb_root *root)
83 {
84 	struct rb_node *parent, *gparent;
85 
86 	if (root->augment_cb)
87 		root->augment_cb(node);
88 
89 	while ((parent = rb_parent(node)) && rb_is_red(parent))
90 	{
91 		gparent = rb_parent(parent);
92 
93 		if (parent == gparent->rb_left)
94 		{
95 			{
96 				register struct rb_node *uncle = gparent->rb_right;
97 				if (uncle && rb_is_red(uncle))
98 				{
99 					rb_set_black(uncle);
100 					rb_set_black(parent);
101 					rb_set_red(gparent);
102 					node = gparent;
103 					continue;
104 				}
105 			}
106 
107 			if (parent->rb_right == node)
108 			{
109 				register struct rb_node *tmp;
110 				__rb_rotate_left(parent, root);
111 				tmp = parent;
112 				parent = node;
113 				node = tmp;
114 			}
115 
116 			rb_set_black(parent);
117 			rb_set_red(gparent);
118 			__rb_rotate_right(gparent, root);
119 		} else {
120 			{
121 				register struct rb_node *uncle = gparent->rb_left;
122 				if (uncle && rb_is_red(uncle))
123 				{
124 					rb_set_black(uncle);
125 					rb_set_black(parent);
126 					rb_set_red(gparent);
127 					node = gparent;
128 					continue;
129 				}
130 			}
131 
132 			if (parent->rb_left == node)
133 			{
134 				register struct rb_node *tmp;
135 				__rb_rotate_right(parent, root);
136 				tmp = parent;
137 				parent = node;
138 				node = tmp;
139 			}
140 
141 			rb_set_black(parent);
142 			rb_set_red(gparent);
143 			__rb_rotate_left(gparent, root);
144 		}
145 	}
146 
147 	rb_set_black(root->rb_node);
148 }
149 EXPORT_SYMBOL(rb_insert_color);
150 
151 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
152 			     struct rb_root *root)
153 {
154 	struct rb_node *other;
155 
156 	while ((!node || rb_is_black(node)) && node != root->rb_node)
157 	{
158 		if (parent->rb_left == node)
159 		{
160 			other = parent->rb_right;
161 			if (rb_is_red(other))
162 			{
163 				rb_set_black(other);
164 				rb_set_red(parent);
165 				__rb_rotate_left(parent, root);
166 				other = parent->rb_right;
167 			}
168 			if ((!other->rb_left || rb_is_black(other->rb_left)) &&
169 			    (!other->rb_right || rb_is_black(other->rb_right)))
170 			{
171 				rb_set_red(other);
172 				node = parent;
173 				parent = rb_parent(node);
174 			}
175 			else
176 			{
177 				if (!other->rb_right || rb_is_black(other->rb_right))
178 				{
179 					rb_set_black(other->rb_left);
180 					rb_set_red(other);
181 					__rb_rotate_right(other, root);
182 					other = parent->rb_right;
183 				}
184 				rb_set_color(other, rb_color(parent));
185 				rb_set_black(parent);
186 				rb_set_black(other->rb_right);
187 				__rb_rotate_left(parent, root);
188 				node = root->rb_node;
189 				break;
190 			}
191 		}
192 		else
193 		{
194 			other = parent->rb_left;
195 			if (rb_is_red(other))
196 			{
197 				rb_set_black(other);
198 				rb_set_red(parent);
199 				__rb_rotate_right(parent, root);
200 				other = parent->rb_left;
201 			}
202 			if ((!other->rb_left || rb_is_black(other->rb_left)) &&
203 			    (!other->rb_right || rb_is_black(other->rb_right)))
204 			{
205 				rb_set_red(other);
206 				node = parent;
207 				parent = rb_parent(node);
208 			}
209 			else
210 			{
211 				if (!other->rb_left || rb_is_black(other->rb_left))
212 				{
213 					rb_set_black(other->rb_right);
214 					rb_set_red(other);
215 					__rb_rotate_left(other, root);
216 					other = parent->rb_left;
217 				}
218 				rb_set_color(other, rb_color(parent));
219 				rb_set_black(parent);
220 				rb_set_black(other->rb_left);
221 				__rb_rotate_right(parent, root);
222 				node = root->rb_node;
223 				break;
224 			}
225 		}
226 	}
227 	if (node)
228 		rb_set_black(node);
229 }
230 
231 void rb_erase(struct rb_node *node, struct rb_root *root)
232 {
233 	struct rb_node *child, *parent;
234 	int color;
235 
236 	if (!node->rb_left)
237 		child = node->rb_right;
238 	else if (!node->rb_right)
239 		child = node->rb_left;
240 	else
241 	{
242 		struct rb_node *old = node, *left;
243 		int old_parent_cb = 0;
244 		int successor_parent_cb = 0;
245 
246 		node = node->rb_right;
247 		while ((left = node->rb_left) != NULL)
248 			node = left;
249 
250 		if (rb_parent(old)) {
251 			old_parent_cb = 1;
252 			if (rb_parent(old)->rb_left == old)
253 				rb_parent(old)->rb_left = node;
254 			else
255 				rb_parent(old)->rb_right = node;
256 		} else
257 			root->rb_node = node;
258 
259 		child = node->rb_right;
260 		parent = rb_parent(node);
261 		color = rb_color(node);
262 
263 		if (parent == old) {
264 			parent = node;
265 		} else {
266 			successor_parent_cb = 1;
267 			if (child)
268 				rb_set_parent(child, parent);
269 
270 			parent->rb_left = child;
271 
272 			node->rb_right = old->rb_right;
273 			rb_set_parent(old->rb_right, node);
274 		}
275 
276 		node->rb_parent_color = old->rb_parent_color;
277 		node->rb_left = old->rb_left;
278 		rb_set_parent(old->rb_left, node);
279 
280 		if (root->augment_cb) {
281 			/*
282 			 * Here, three different nodes can have new children.
283 			 * The parent of the successor node that was selected
284 			 * to replace the node to be erased.
285 			 * The node that is getting erased and is now replaced
286 			 * by its successor.
287 			 * The parent of the node getting erased-replaced.
288 			 */
289 			if (successor_parent_cb)
290 				root->augment_cb(parent);
291 
292 			root->augment_cb(node);
293 
294 			if (old_parent_cb)
295 				root->augment_cb(rb_parent(old));
296 		}
297 
298 		goto color;
299 	}
300 
301 	parent = rb_parent(node);
302 	color = rb_color(node);
303 
304 	if (child)
305 		rb_set_parent(child, parent);
306 
307 	if (parent) {
308 		if (parent->rb_left == node)
309 			parent->rb_left = child;
310 		else
311 			parent->rb_right = child;
312 
313 		if (root->augment_cb)
314 			root->augment_cb(parent);
315 
316 	} else {
317 		root->rb_node = child;
318 	}
319 
320  color:
321 	if (color == RB_BLACK)
322 		__rb_erase_color(child, parent, root);
323 }
324 EXPORT_SYMBOL(rb_erase);
325 
326 /*
327  * This function returns the first node (in sort order) of the tree.
328  */
329 struct rb_node *rb_first(const struct rb_root *root)
330 {
331 	struct rb_node	*n;
332 
333 	n = root->rb_node;
334 	if (!n)
335 		return NULL;
336 	while (n->rb_left)
337 		n = n->rb_left;
338 	return n;
339 }
340 EXPORT_SYMBOL(rb_first);
341 
342 struct rb_node *rb_last(const struct rb_root *root)
343 {
344 	struct rb_node	*n;
345 
346 	n = root->rb_node;
347 	if (!n)
348 		return NULL;
349 	while (n->rb_right)
350 		n = n->rb_right;
351 	return n;
352 }
353 EXPORT_SYMBOL(rb_last);
354 
355 struct rb_node *rb_next(const struct rb_node *node)
356 {
357 	struct rb_node *parent;
358 
359 	if (rb_parent(node) == node)
360 		return NULL;
361 
362 	/* If we have a right-hand child, go down and then left as far
363 	   as we can. */
364 	if (node->rb_right) {
365 		node = node->rb_right;
366 		while (node->rb_left)
367 			node=node->rb_left;
368 		return (struct rb_node *)node;
369 	}
370 
371 	/* No right-hand children.  Everything down and left is
372 	   smaller than us, so any 'next' node must be in the general
373 	   direction of our parent. Go up the tree; any time the
374 	   ancestor is a right-hand child of its parent, keep going
375 	   up. First time it's a left-hand child of its parent, said
376 	   parent is our 'next' node. */
377 	while ((parent = rb_parent(node)) && node == parent->rb_right)
378 		node = parent;
379 
380 	return parent;
381 }
382 EXPORT_SYMBOL(rb_next);
383 
384 struct rb_node *rb_prev(const struct rb_node *node)
385 {
386 	struct rb_node *parent;
387 
388 	if (rb_parent(node) == node)
389 		return NULL;
390 
391 	/* If we have a left-hand child, go down and then right as far
392 	   as we can. */
393 	if (node->rb_left) {
394 		node = node->rb_left;
395 		while (node->rb_right)
396 			node=node->rb_right;
397 		return (struct rb_node *)node;
398 	}
399 
400 	/* No left-hand children. Go up till we find an ancestor which
401 	   is a right-hand child of its parent */
402 	while ((parent = rb_parent(node)) && node == parent->rb_left)
403 		node = parent;
404 
405 	return parent;
406 }
407 EXPORT_SYMBOL(rb_prev);
408 
409 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
410 		     struct rb_root *root)
411 {
412 	struct rb_node *parent = rb_parent(victim);
413 
414 	/* Set the surrounding nodes to point to the replacement */
415 	if (parent) {
416 		if (victim == parent->rb_left)
417 			parent->rb_left = new;
418 		else
419 			parent->rb_right = new;
420 	} else {
421 		root->rb_node = new;
422 	}
423 	if (victim->rb_left)
424 		rb_set_parent(victim->rb_left, new);
425 	if (victim->rb_right)
426 		rb_set_parent(victim->rb_right, new);
427 
428 	/* Copy the pointers/colour from the victim to the replacement */
429 	*new = *victim;
430 }
431 EXPORT_SYMBOL(rb_replace_node);
432