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 "rbtree.h"
24 
__rb_rotate_left(struct rb_node * node,struct rb_root * root)25 static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
26 {
27 	struct rb_node *right = node->rb_right;
28 	struct rb_node *parent = ext2fs_rb_parent(node);
29 
30 	if ((node->rb_right = right->rb_left))
31 		ext2fs_rb_set_parent(right->rb_left, node);
32 	right->rb_left = node;
33 
34 	ext2fs_rb_set_parent(right, parent);
35 
36 	if (parent)
37 	{
38 		if (node == parent->rb_left)
39 			parent->rb_left = right;
40 		else
41 			parent->rb_right = right;
42 	}
43 	else
44 		root->rb_node = right;
45 	ext2fs_rb_set_parent(node, right);
46 }
47 
__rb_rotate_right(struct rb_node * node,struct rb_root * root)48 static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
49 {
50 	struct rb_node *left = node->rb_left;
51 	struct rb_node *parent = ext2fs_rb_parent(node);
52 
53 	if ((node->rb_left = left->rb_right))
54 		ext2fs_rb_set_parent(left->rb_right, node);
55 	left->rb_right = node;
56 
57 	ext2fs_rb_set_parent(left, parent);
58 
59 	if (parent)
60 	{
61 		if (node == parent->rb_right)
62 			parent->rb_right = left;
63 		else
64 			parent->rb_left = left;
65 	}
66 	else
67 		root->rb_node = left;
68 	ext2fs_rb_set_parent(node, left);
69 }
70 
ext2fs_rb_insert_color(struct rb_node * node,struct rb_root * root)71 void ext2fs_rb_insert_color(struct rb_node *node, struct rb_root *root)
72 {
73 	struct rb_node *parent, *gparent;
74 
75 	while ((parent = ext2fs_rb_parent(node)) && ext2fs_rb_is_red(parent))
76 	{
77 		gparent = ext2fs_rb_parent(parent);
78 
79 		if (parent == gparent->rb_left)
80 		{
81 			{
82 				register struct rb_node *uncle = gparent->rb_right;
83 				if (uncle && ext2fs_rb_is_red(uncle))
84 				{
85 					ext2fs_rb_set_black(uncle);
86 					ext2fs_rb_set_black(parent);
87 					ext2fs_rb_set_red(gparent);
88 					node = gparent;
89 					continue;
90 				}
91 			}
92 
93 			if (parent->rb_right == node)
94 			{
95 				register struct rb_node *tmp;
96 				__rb_rotate_left(parent, root);
97 				tmp = parent;
98 				parent = node;
99 				node = tmp;
100 			}
101 
102 			ext2fs_rb_set_black(parent);
103 			ext2fs_rb_set_red(gparent);
104 			__rb_rotate_right(gparent, root);
105 		} else {
106 			{
107 				register struct rb_node *uncle = gparent->rb_left;
108 				if (uncle && ext2fs_rb_is_red(uncle))
109 				{
110 					ext2fs_rb_set_black(uncle);
111 					ext2fs_rb_set_black(parent);
112 					ext2fs_rb_set_red(gparent);
113 					node = gparent;
114 					continue;
115 				}
116 			}
117 
118 			if (parent->rb_left == node)
119 			{
120 				register struct rb_node *tmp;
121 				__rb_rotate_right(parent, root);
122 				tmp = parent;
123 				parent = node;
124 				node = tmp;
125 			}
126 
127 			ext2fs_rb_set_black(parent);
128 			ext2fs_rb_set_red(gparent);
129 			__rb_rotate_left(gparent, root);
130 		}
131 	}
132 
133 	ext2fs_rb_set_black(root->rb_node);
134 }
135 
__rb_erase_color(struct rb_node * node,struct rb_node * parent,struct rb_root * root)136 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
137 			     struct rb_root *root)
138 {
139 	struct rb_node *other;
140 
141 	while ((!node || ext2fs_rb_is_black(node)) && node != root->rb_node)
142 	{
143 		if (parent->rb_left == node)
144 		{
145 			other = parent->rb_right;
146 			if (ext2fs_rb_is_red(other))
147 			{
148 				ext2fs_rb_set_black(other);
149 				ext2fs_rb_set_red(parent);
150 				__rb_rotate_left(parent, root);
151 				other = parent->rb_right;
152 			}
153 			if ((!other->rb_left || ext2fs_rb_is_black(other->rb_left)) &&
154 			    (!other->rb_right || ext2fs_rb_is_black(other->rb_right)))
155 			{
156 				ext2fs_rb_set_red(other);
157 				node = parent;
158 				parent = ext2fs_rb_parent(node);
159 			}
160 			else
161 			{
162 				if (!other->rb_right || ext2fs_rb_is_black(other->rb_right))
163 				{
164 					ext2fs_rb_set_black(other->rb_left);
165 					ext2fs_rb_set_red(other);
166 					__rb_rotate_right(other, root);
167 					other = parent->rb_right;
168 				}
169 				ext2fs_rb_set_color(other, ext2fs_rb_color(parent));
170 				ext2fs_rb_set_black(parent);
171 				ext2fs_rb_set_black(other->rb_right);
172 				__rb_rotate_left(parent, root);
173 				node = root->rb_node;
174 				break;
175 			}
176 		}
177 		else
178 		{
179 			other = parent->rb_left;
180 			if (ext2fs_rb_is_red(other))
181 			{
182 				ext2fs_rb_set_black(other);
183 				ext2fs_rb_set_red(parent);
184 				__rb_rotate_right(parent, root);
185 				other = parent->rb_left;
186 			}
187 			if ((!other->rb_left || ext2fs_rb_is_black(other->rb_left)) &&
188 			    (!other->rb_right || ext2fs_rb_is_black(other->rb_right)))
189 			{
190 				ext2fs_rb_set_red(other);
191 				node = parent;
192 				parent = ext2fs_rb_parent(node);
193 			}
194 			else
195 			{
196 				if (!other->rb_left || ext2fs_rb_is_black(other->rb_left))
197 				{
198 					ext2fs_rb_set_black(other->rb_right);
199 					ext2fs_rb_set_red(other);
200 					__rb_rotate_left(other, root);
201 					other = parent->rb_left;
202 				}
203 				ext2fs_rb_set_color(other, ext2fs_rb_color(parent));
204 				ext2fs_rb_set_black(parent);
205 				ext2fs_rb_set_black(other->rb_left);
206 				__rb_rotate_right(parent, root);
207 				node = root->rb_node;
208 				break;
209 			}
210 		}
211 	}
212 	if (node)
213 		ext2fs_rb_set_black(node);
214 }
215 
ext2fs_rb_erase(struct rb_node * node,struct rb_root * root)216 void ext2fs_rb_erase(struct rb_node *node, struct rb_root *root)
217 {
218 	struct rb_node *child, *parent;
219 	int color;
220 
221 	if (!node->rb_left)
222 		child = node->rb_right;
223 	else if (!node->rb_right)
224 		child = node->rb_left;
225 	else
226 	{
227 		struct rb_node *old = node, *left;
228 
229 		node = node->rb_right;
230 		while ((left = node->rb_left) != NULL)
231 			node = left;
232 
233 		if (ext2fs_rb_parent(old)) {
234 			if (ext2fs_rb_parent(old)->rb_left == old)
235 				ext2fs_rb_parent(old)->rb_left = node;
236 			else
237 				ext2fs_rb_parent(old)->rb_right = node;
238 		} else
239 			root->rb_node = node;
240 
241 		child = node->rb_right;
242 		parent = ext2fs_rb_parent(node);
243 		color = ext2fs_rb_color(node);
244 
245 		if (parent == old) {
246 			parent = node;
247 		} else {
248 			if (child)
249 				ext2fs_rb_set_parent(child, parent);
250 			parent->rb_left = child;
251 
252 			node->rb_right = old->rb_right;
253 			ext2fs_rb_set_parent(old->rb_right, node);
254 		}
255 
256 		node->rb_parent_color = old->rb_parent_color;
257 		node->rb_left = old->rb_left;
258 		ext2fs_rb_set_parent(old->rb_left, node);
259 
260 		goto color;
261 	}
262 
263 	parent = ext2fs_rb_parent(node);
264 	color = ext2fs_rb_color(node);
265 
266 	if (child)
267 		ext2fs_rb_set_parent(child, parent);
268 	if (parent)
269 	{
270 		if (parent->rb_left == node)
271 			parent->rb_left = child;
272 		else
273 			parent->rb_right = child;
274 	}
275 	else
276 		root->rb_node = child;
277 
278  color:
279 	if (color == RB_BLACK)
280 		__rb_erase_color(child, parent, root);
281 }
282 
ext2fs_rb_augment_path(struct rb_node * node,rb_augment_f func,void * data)283 static void ext2fs_rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
284 {
285 	struct rb_node *parent;
286 
287 up:
288 	func(node, data);
289 	parent = ext2fs_rb_parent(node);
290 	if (!parent)
291 		return;
292 
293 	if (node == parent->rb_left && parent->rb_right)
294 		func(parent->rb_right, data);
295 	else if (parent->rb_left)
296 		func(parent->rb_left, data);
297 
298 	node = parent;
299 	goto up;
300 }
301 
302 /*
303  * after inserting @node into the tree, update the tree to account for
304  * both the new entry and any damage done by rebalance
305  */
ext2fs_rb_augment_insert(struct rb_node * node,rb_augment_f func,void * data)306 void ext2fs_rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
307 {
308 	if (node->rb_left)
309 		node = node->rb_left;
310 	else if (node->rb_right)
311 		node = node->rb_right;
312 
313 	ext2fs_rb_augment_path(node, func, data);
314 }
315 
316 /*
317  * before removing the node, find the deepest node on the rebalance path
318  * that will still be there after @node gets removed
319  */
ext2fs_rb_augment_erase_begin(struct rb_node * node)320 struct rb_node *ext2fs_rb_augment_erase_begin(struct rb_node *node)
321 {
322 	struct rb_node *deepest;
323 
324 	if (!node->rb_right && !node->rb_left)
325 		deepest = ext2fs_rb_parent(node);
326 	else if (!node->rb_right)
327 		deepest = node->rb_left;
328 	else if (!node->rb_left)
329 		deepest = node->rb_right;
330 	else {
331 		deepest = ext2fs_rb_next(node);
332 		if (deepest->rb_right)
333 			deepest = deepest->rb_right;
334 		else if (ext2fs_rb_parent(deepest) != node)
335 			deepest = ext2fs_rb_parent(deepest);
336 	}
337 
338 	return deepest;
339 }
340 
341 /*
342  * after removal, update the tree to account for the removed entry
343  * and any rebalance damage.
344  */
ext2fs_rb_augment_erase_end(struct rb_node * node,rb_augment_f func,void * data)345 void ext2fs_rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
346 {
347 	if (node)
348 		ext2fs_rb_augment_path(node, func, data);
349 }
350 
351 /*
352  * This function returns the first node (in sort order) of the tree.
353  */
ext2fs_rb_first(const struct rb_root * root)354 struct rb_node *ext2fs_rb_first(const struct rb_root *root)
355 {
356 	struct rb_node	*n;
357 
358 	n = root->rb_node;
359 	if (!n)
360 		return NULL;
361 	while (n->rb_left)
362 		n = n->rb_left;
363 	return n;
364 }
365 
ext2fs_rb_last(const struct rb_root * root)366 struct rb_node *ext2fs_rb_last(const struct rb_root *root)
367 {
368 	struct rb_node	*n;
369 
370 	n = root->rb_node;
371 	if (!n)
372 		return NULL;
373 	while (n->rb_right)
374 		n = n->rb_right;
375 	return n;
376 }
377 
ext2fs_rb_next(struct rb_node * node)378 struct rb_node *ext2fs_rb_next(struct rb_node *node)
379 {
380 	struct rb_node *parent;
381 
382 	if (ext2fs_rb_parent(node) == node)
383 		return NULL;
384 
385 	/* If we have a right-hand child, go down and then left as far
386 	   as we can. */
387 	if (node->rb_right) {
388 		node = node->rb_right;
389 		while (node->rb_left)
390 			node=node->rb_left;
391 		return (struct rb_node *)node;
392 	}
393 
394 	/* No right-hand children.  Everything down and left is
395 	   smaller than us, so any 'next' node must be in the general
396 	   direction of our parent. Go up the tree; any time the
397 	   ancestor is a right-hand child of its parent, keep going
398 	   up. First time it's a left-hand child of its parent, said
399 	   parent is our 'next' node. */
400 	while ((parent = ext2fs_rb_parent(node)) && node == parent->rb_right)
401 		node = parent;
402 
403 	return parent;
404 }
405 
ext2fs_rb_prev(struct rb_node * node)406 struct rb_node *ext2fs_rb_prev(struct rb_node *node)
407 {
408 	struct rb_node *parent;
409 
410 	if (ext2fs_rb_parent(node) == node)
411 		return NULL;
412 
413 	/* If we have a left-hand child, go down and then right as far
414 	   as we can. */
415 	if (node->rb_left) {
416 		node = node->rb_left;
417 		while (node->rb_right)
418 			node=node->rb_right;
419 		return (struct rb_node *)node;
420 	}
421 
422 	/* No left-hand children. Go up till we find an ancestor which
423 	   is a right-hand child of its parent */
424 	while ((parent = ext2fs_rb_parent(node)) && node == parent->rb_left)
425 		node = parent;
426 
427 	return parent;
428 }
429 
ext2fs_rb_replace_node(struct rb_node * victim,struct rb_node * new,struct rb_root * root)430 void ext2fs_rb_replace_node(struct rb_node *victim, struct rb_node *new,
431 			  struct rb_root *root)
432 {
433 	struct rb_node *parent = ext2fs_rb_parent(victim);
434 
435 	/* Set the surrounding nodes to point to the replacement */
436 	if (parent) {
437 		if (victim == parent->rb_left)
438 			parent->rb_left = new;
439 		else
440 			parent->rb_right = new;
441 	} else {
442 		root->rb_node = new;
443 	}
444 	if (victim->rb_left)
445 		ext2fs_rb_set_parent(victim->rb_left, new);
446 	if (victim->rb_right)
447 		ext2fs_rb_set_parent(victim->rb_right, new);
448 
449 	/* Copy the pointers/colour from the victim to the replacement */
450 	*new = *victim;
451 }
452