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