xref: /openbmc/u-boot/lib/rbtree.c (revision d9b88d25)
1 /*
2   Red Black Trees
3   (C) 1999  Andrea Arcangeli <andrea@suse.de>
4   (C) 2002  David Woodhouse <dwmw2@infradead.org>
5   (C) 2012  Michel Lespinasse <walken@google.com>
6 
7  * SPDX-License-Identifier:	GPL-2.0+
8 
9   linux/lib/rbtree.c
10 */
11 
12 #include <linux/rbtree_augmented.h>
13 #ifndef __UBOOT__
14 #include <linux/export.h>
15 #else
16 #include <ubi_uboot.h>
17 #endif
18 /*
19  * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
20  *
21  *  1) A node is either red or black
22  *  2) The root is black
23  *  3) All leaves (NULL) are black
24  *  4) Both children of every red node are black
25  *  5) Every simple path from root to leaves contains the same number
26  *     of black nodes.
27  *
28  *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
29  *  consecutive red nodes in a path and every red node is therefore followed by
30  *  a black. So if B is the number of black nodes on every simple path (as per
31  *  5), then the longest possible path due to 4 is 2B.
32  *
33  *  We shall indicate color with case, where black nodes are uppercase and red
34  *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
35  *  parentheses and have some accompanying text comment.
36  */
37 
38 static inline void rb_set_black(struct rb_node *rb)
39 {
40 	rb->__rb_parent_color |= RB_BLACK;
41 }
42 
43 static inline struct rb_node *rb_red_parent(struct rb_node *red)
44 {
45 	return (struct rb_node *)red->__rb_parent_color;
46 }
47 
48 /*
49  * Helper function for rotations:
50  * - old's parent and color get assigned to new
51  * - old gets assigned new as a parent and 'color' as a color.
52  */
53 static inline void
54 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
55 			struct rb_root *root, int color)
56 {
57 	struct rb_node *parent = rb_parent(old);
58 	new->__rb_parent_color = old->__rb_parent_color;
59 	rb_set_parent_color(old, new, color);
60 	__rb_change_child(old, new, parent, root);
61 }
62 
63 static __always_inline void
64 __rb_insert(struct rb_node *node, struct rb_root *root,
65 	    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
66 {
67 	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
68 
69 	while (true) {
70 		/*
71 		 * Loop invariant: node is red
72 		 *
73 		 * If there is a black parent, we are done.
74 		 * Otherwise, take some corrective action as we don't
75 		 * want a red root or two consecutive red nodes.
76 		 */
77 		if (!parent) {
78 			rb_set_parent_color(node, NULL, RB_BLACK);
79 			break;
80 		} else if (rb_is_black(parent))
81 			break;
82 
83 		gparent = rb_red_parent(parent);
84 
85 		tmp = gparent->rb_right;
86 		if (parent != tmp) {	/* parent == gparent->rb_left */
87 			if (tmp && rb_is_red(tmp)) {
88 				/*
89 				 * Case 1 - color flips
90 				 *
91 				 *       G            g
92 				 *      / \          / \
93 				 *     p   u  -->   P   U
94 				 *    /            /
95 				 *   n            N
96 				 *
97 				 * However, since g's parent might be red, and
98 				 * 4) does not allow this, we need to recurse
99 				 * at g.
100 				 */
101 				rb_set_parent_color(tmp, gparent, RB_BLACK);
102 				rb_set_parent_color(parent, gparent, RB_BLACK);
103 				node = gparent;
104 				parent = rb_parent(node);
105 				rb_set_parent_color(node, parent, RB_RED);
106 				continue;
107 			}
108 
109 			tmp = parent->rb_right;
110 			if (node == tmp) {
111 				/*
112 				 * Case 2 - left rotate at parent
113 				 *
114 				 *      G             G
115 				 *     / \           / \
116 				 *    p   U  -->    n   U
117 				 *     \           /
118 				 *      n         p
119 				 *
120 				 * This still leaves us in violation of 4), the
121 				 * continuation into Case 3 will fix that.
122 				 */
123 				parent->rb_right = tmp = node->rb_left;
124 				node->rb_left = parent;
125 				if (tmp)
126 					rb_set_parent_color(tmp, parent,
127 							    RB_BLACK);
128 				rb_set_parent_color(parent, node, RB_RED);
129 				augment_rotate(parent, node);
130 				parent = node;
131 				tmp = node->rb_right;
132 			}
133 
134 			/*
135 			 * Case 3 - right rotate at gparent
136 			 *
137 			 *        G           P
138 			 *       / \         / \
139 			 *      p   U  -->  n   g
140 			 *     /                 \
141 			 *    n                   U
142 			 */
143 			gparent->rb_left = tmp;  /* == parent->rb_right */
144 			parent->rb_right = gparent;
145 			if (tmp)
146 				rb_set_parent_color(tmp, gparent, RB_BLACK);
147 			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
148 			augment_rotate(gparent, parent);
149 			break;
150 		} else {
151 			tmp = gparent->rb_left;
152 			if (tmp && rb_is_red(tmp)) {
153 				/* Case 1 - color flips */
154 				rb_set_parent_color(tmp, gparent, RB_BLACK);
155 				rb_set_parent_color(parent, gparent, RB_BLACK);
156 				node = gparent;
157 				parent = rb_parent(node);
158 				rb_set_parent_color(node, parent, RB_RED);
159 				continue;
160 			}
161 
162 			tmp = parent->rb_left;
163 			if (node == tmp) {
164 				/* Case 2 - right rotate at parent */
165 				parent->rb_left = tmp = node->rb_right;
166 				node->rb_right = parent;
167 				if (tmp)
168 					rb_set_parent_color(tmp, parent,
169 							    RB_BLACK);
170 				rb_set_parent_color(parent, node, RB_RED);
171 				augment_rotate(parent, node);
172 				parent = node;
173 				tmp = node->rb_left;
174 			}
175 
176 			/* Case 3 - left rotate at gparent */
177 			gparent->rb_right = tmp;  /* == parent->rb_left */
178 			parent->rb_left = gparent;
179 			if (tmp)
180 				rb_set_parent_color(tmp, gparent, RB_BLACK);
181 			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
182 			augment_rotate(gparent, parent);
183 			break;
184 		}
185 	}
186 }
187 
188 /*
189  * Inline version for rb_erase() use - we want to be able to inline
190  * and eliminate the dummy_rotate callback there
191  */
192 static __always_inline void
193 ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
194 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
195 {
196 	struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
197 
198 	while (true) {
199 		/*
200 		 * Loop invariants:
201 		 * - node is black (or NULL on first iteration)
202 		 * - node is not the root (parent is not NULL)
203 		 * - All leaf paths going through parent and node have a
204 		 *   black node count that is 1 lower than other leaf paths.
205 		 */
206 		sibling = parent->rb_right;
207 		if (node != sibling) {	/* node == parent->rb_left */
208 			if (rb_is_red(sibling)) {
209 				/*
210 				 * Case 1 - left rotate at parent
211 				 *
212 				 *     P               S
213 				 *    / \             / \
214 				 *   N   s    -->    p   Sr
215 				 *      / \         / \
216 				 *     Sl  Sr      N   Sl
217 				 */
218 				parent->rb_right = tmp1 = sibling->rb_left;
219 				sibling->rb_left = parent;
220 				rb_set_parent_color(tmp1, parent, RB_BLACK);
221 				__rb_rotate_set_parents(parent, sibling, root,
222 							RB_RED);
223 				augment_rotate(parent, sibling);
224 				sibling = tmp1;
225 			}
226 			tmp1 = sibling->rb_right;
227 			if (!tmp1 || rb_is_black(tmp1)) {
228 				tmp2 = sibling->rb_left;
229 				if (!tmp2 || rb_is_black(tmp2)) {
230 					/*
231 					 * Case 2 - sibling color flip
232 					 * (p could be either color here)
233 					 *
234 					 *    (p)           (p)
235 					 *    / \           / \
236 					 *   N   S    -->  N   s
237 					 *      / \           / \
238 					 *     Sl  Sr        Sl  Sr
239 					 *
240 					 * This leaves us violating 5) which
241 					 * can be fixed by flipping p to black
242 					 * if it was red, or by recursing at p.
243 					 * p is red when coming from Case 1.
244 					 */
245 					rb_set_parent_color(sibling, parent,
246 							    RB_RED);
247 					if (rb_is_red(parent))
248 						rb_set_black(parent);
249 					else {
250 						node = parent;
251 						parent = rb_parent(node);
252 						if (parent)
253 							continue;
254 					}
255 					break;
256 				}
257 				/*
258 				 * Case 3 - right rotate at sibling
259 				 * (p could be either color here)
260 				 *
261 				 *   (p)           (p)
262 				 *   / \           / \
263 				 *  N   S    -->  N   Sl
264 				 *     / \             \
265 				 *    sl  Sr            s
266 				 *                       \
267 				 *                        Sr
268 				 */
269 				sibling->rb_left = tmp1 = tmp2->rb_right;
270 				tmp2->rb_right = sibling;
271 				parent->rb_right = tmp2;
272 				if (tmp1)
273 					rb_set_parent_color(tmp1, sibling,
274 							    RB_BLACK);
275 				augment_rotate(sibling, tmp2);
276 				tmp1 = sibling;
277 				sibling = tmp2;
278 			}
279 			/*
280 			 * Case 4 - left rotate at parent + color flips
281 			 * (p and sl could be either color here.
282 			 *  After rotation, p becomes black, s acquires
283 			 *  p's color, and sl keeps its color)
284 			 *
285 			 *      (p)             (s)
286 			 *      / \             / \
287 			 *     N   S     -->   P   Sr
288 			 *        / \         / \
289 			 *      (sl) sr      N  (sl)
290 			 */
291 			parent->rb_right = tmp2 = sibling->rb_left;
292 			sibling->rb_left = parent;
293 			rb_set_parent_color(tmp1, sibling, RB_BLACK);
294 			if (tmp2)
295 				rb_set_parent(tmp2, parent);
296 			__rb_rotate_set_parents(parent, sibling, root,
297 						RB_BLACK);
298 			augment_rotate(parent, sibling);
299 			break;
300 		} else {
301 			sibling = parent->rb_left;
302 			if (rb_is_red(sibling)) {
303 				/* Case 1 - right rotate at parent */
304 				parent->rb_left = tmp1 = sibling->rb_right;
305 				sibling->rb_right = parent;
306 				rb_set_parent_color(tmp1, parent, RB_BLACK);
307 				__rb_rotate_set_parents(parent, sibling, root,
308 							RB_RED);
309 				augment_rotate(parent, sibling);
310 				sibling = tmp1;
311 			}
312 			tmp1 = sibling->rb_left;
313 			if (!tmp1 || rb_is_black(tmp1)) {
314 				tmp2 = sibling->rb_right;
315 				if (!tmp2 || rb_is_black(tmp2)) {
316 					/* Case 2 - sibling color flip */
317 					rb_set_parent_color(sibling, parent,
318 							    RB_RED);
319 					if (rb_is_red(parent))
320 						rb_set_black(parent);
321 					else {
322 						node = parent;
323 						parent = rb_parent(node);
324 						if (parent)
325 							continue;
326 					}
327 					break;
328 				}
329 				/* Case 3 - right rotate at sibling */
330 				sibling->rb_right = tmp1 = tmp2->rb_left;
331 				tmp2->rb_left = sibling;
332 				parent->rb_left = tmp2;
333 				if (tmp1)
334 					rb_set_parent_color(tmp1, sibling,
335 							    RB_BLACK);
336 				augment_rotate(sibling, tmp2);
337 				tmp1 = sibling;
338 				sibling = tmp2;
339 			}
340 			/* Case 4 - left rotate at parent + color flips */
341 			parent->rb_left = tmp2 = sibling->rb_right;
342 			sibling->rb_right = parent;
343 			rb_set_parent_color(tmp1, sibling, RB_BLACK);
344 			if (tmp2)
345 				rb_set_parent(tmp2, parent);
346 			__rb_rotate_set_parents(parent, sibling, root,
347 						RB_BLACK);
348 			augment_rotate(parent, sibling);
349 			break;
350 		}
351 	}
352 }
353 
354 /* Non-inline version for rb_erase_augmented() use */
355 void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
356 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
357 {
358 	____rb_erase_color(parent, root, augment_rotate);
359 }
360 EXPORT_SYMBOL(__rb_erase_color);
361 
362 /*
363  * Non-augmented rbtree manipulation functions.
364  *
365  * We use dummy augmented callbacks here, and have the compiler optimize them
366  * out of the rb_insert_color() and rb_erase() function definitions.
367  */
368 
369 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
370 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
371 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
372 
373 static const struct rb_augment_callbacks dummy_callbacks = {
374 	dummy_propagate, dummy_copy, dummy_rotate
375 };
376 
377 void rb_insert_color(struct rb_node *node, struct rb_root *root)
378 {
379 	__rb_insert(node, root, dummy_rotate);
380 }
381 EXPORT_SYMBOL(rb_insert_color);
382 
383 void rb_erase(struct rb_node *node, struct rb_root *root)
384 {
385 	struct rb_node *rebalance;
386 	rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
387 	if (rebalance)
388 		____rb_erase_color(rebalance, root, dummy_rotate);
389 }
390 EXPORT_SYMBOL(rb_erase);
391 
392 /*
393  * Augmented rbtree manipulation functions.
394  *
395  * This instantiates the same __always_inline functions as in the non-augmented
396  * case, but this time with user-defined callbacks.
397  */
398 
399 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
400 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
401 {
402 	__rb_insert(node, root, augment_rotate);
403 }
404 EXPORT_SYMBOL(__rb_insert_augmented);
405 
406 /*
407  * This function returns the first node (in sort order) of the tree.
408  */
409 struct rb_node *rb_first(const struct rb_root *root)
410 {
411 	struct rb_node	*n;
412 
413 	n = root->rb_node;
414 	if (!n)
415 		return NULL;
416 	while (n->rb_left)
417 		n = n->rb_left;
418 	return n;
419 }
420 EXPORT_SYMBOL(rb_first);
421 
422 struct rb_node *rb_last(const struct rb_root *root)
423 {
424 	struct rb_node	*n;
425 
426 	n = root->rb_node;
427 	if (!n)
428 		return NULL;
429 	while (n->rb_right)
430 		n = n->rb_right;
431 	return n;
432 }
433 EXPORT_SYMBOL(rb_last);
434 
435 struct rb_node *rb_next(const struct rb_node *node)
436 {
437 	struct rb_node *parent;
438 
439 	if (RB_EMPTY_NODE(node))
440 		return NULL;
441 
442 	/*
443 	 * If we have a right-hand child, go down and then left as far
444 	 * as we can.
445 	 */
446 	if (node->rb_right) {
447 		node = node->rb_right;
448 		while (node->rb_left)
449 			node=node->rb_left;
450 		return (struct rb_node *)node;
451 	}
452 
453 	/*
454 	 * No right-hand children. Everything down and left is smaller than us,
455 	 * so any 'next' node must be in the general direction of our parent.
456 	 * Go up the tree; any time the ancestor is a right-hand child of its
457 	 * parent, keep going up. First time it's a left-hand child of its
458 	 * parent, said parent is our 'next' node.
459 	 */
460 	while ((parent = rb_parent(node)) && node == parent->rb_right)
461 		node = parent;
462 
463 	return parent;
464 }
465 EXPORT_SYMBOL(rb_next);
466 
467 struct rb_node *rb_prev(const struct rb_node *node)
468 {
469 	struct rb_node *parent;
470 
471 	if (RB_EMPTY_NODE(node))
472 		return NULL;
473 
474 	/*
475 	 * If we have a left-hand child, go down and then right as far
476 	 * as we can.
477 	 */
478 	if (node->rb_left) {
479 		node = node->rb_left;
480 		while (node->rb_right)
481 			node=node->rb_right;
482 		return (struct rb_node *)node;
483 	}
484 
485 	/*
486 	 * No left-hand children. Go up till we find an ancestor which
487 	 * is a right-hand child of its parent.
488 	 */
489 	while ((parent = rb_parent(node)) && node == parent->rb_left)
490 		node = parent;
491 
492 	return parent;
493 }
494 EXPORT_SYMBOL(rb_prev);
495 
496 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
497 		     struct rb_root *root)
498 {
499 	struct rb_node *parent = rb_parent(victim);
500 
501 	/* Set the surrounding nodes to point to the replacement */
502 	__rb_change_child(victim, new, parent, root);
503 	if (victim->rb_left)
504 		rb_set_parent(victim->rb_left, new);
505 	if (victim->rb_right)
506 		rb_set_parent(victim->rb_right, new);
507 
508 	/* Copy the pointers/colour from the victim to the replacement */
509 	*new = *victim;
510 }
511 EXPORT_SYMBOL(rb_replace_node);
512 
513 static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
514 {
515 	for (;;) {
516 		if (node->rb_left)
517 			node = node->rb_left;
518 		else if (node->rb_right)
519 			node = node->rb_right;
520 		else
521 			return (struct rb_node *)node;
522 	}
523 }
524 
525 struct rb_node *rb_next_postorder(const struct rb_node *node)
526 {
527 	const struct rb_node *parent;
528 	if (!node)
529 		return NULL;
530 	parent = rb_parent(node);
531 
532 	/* If we're sitting on node, we've already seen our children */
533 	if (parent && node == parent->rb_left && parent->rb_right) {
534 		/* If we are the parent's left node, go to the parent's right
535 		 * node then all the way down to the left */
536 		return rb_left_deepest_node(parent->rb_right);
537 	} else
538 		/* Otherwise we are the parent's right node, and the parent
539 		 * should be next */
540 		return (struct rb_node *)parent;
541 }
542 EXPORT_SYMBOL(rb_next_postorder);
543 
544 struct rb_node *rb_first_postorder(const struct rb_root *root)
545 {
546 	if (!root->rb_node)
547 		return NULL;
548 
549 	return rb_left_deepest_node(root->rb_node);
550 }
551 EXPORT_SYMBOL(rb_first_postorder);
552