Merge tag 'arm64-upstream' of git://git.kernel.org/pub/scm/linux/kernel/git/arm64...
[platform/kernel/linux-rpi.git] / lib / rbtree.c
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   This program is free software; you can redistribute it and/or modify
8   it under the terms of the GNU General Public License as published by
9   the Free Software Foundation; either version 2 of the License, or
10   (at your option) any later version.
11
12   This program is distributed in the hope that it will be useful,
13   but WITHOUT ANY WARRANTY; without even the implied warranty of
14   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15   GNU General Public License for more details.
16
17   You should have received a copy of the GNU General Public License
18   along with this program; if not, write to the Free Software
19   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
20
21   linux/lib/rbtree.c
22 */
23
24 #include <linux/rbtree_augmented.h>
25 #include <linux/export.h>
26
27 /*
28  * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
29  *
30  *  1) A node is either red or black
31  *  2) The root is black
32  *  3) All leaves (NULL) are black
33  *  4) Both children of every red node are black
34  *  5) Every simple path from root to leaves contains the same number
35  *     of black nodes.
36  *
37  *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
38  *  consecutive red nodes in a path and every red node is therefore followed by
39  *  a black. So if B is the number of black nodes on every simple path (as per
40  *  5), then the longest possible path due to 4 is 2B.
41  *
42  *  We shall indicate color with case, where black nodes are uppercase and red
43  *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
44  *  parentheses and have some accompanying text comment.
45  */
46
47 /*
48  * Notes on lockless lookups:
49  *
50  * All stores to the tree structure (rb_left and rb_right) must be done using
51  * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
52  * tree structure as seen in program order.
53  *
54  * These two requirements will allow lockless iteration of the tree -- not
55  * correct iteration mind you, tree rotations are not atomic so a lookup might
56  * miss entire subtrees.
57  *
58  * But they do guarantee that any such traversal will only see valid elements
59  * and that it will indeed complete -- does not get stuck in a loop.
60  *
61  * It also guarantees that if the lookup returns an element it is the 'correct'
62  * one. But not returning an element does _NOT_ mean it's not present.
63  *
64  * NOTE:
65  *
66  * Stores to __rb_parent_color are not important for simple lookups so those
67  * are left undone as of now. Nor did I check for loops involving parent
68  * pointers.
69  */
70
71 static inline void rb_set_black(struct rb_node *rb)
72 {
73         rb->__rb_parent_color |= RB_BLACK;
74 }
75
76 static inline struct rb_node *rb_red_parent(struct rb_node *red)
77 {
78         return (struct rb_node *)red->__rb_parent_color;
79 }
80
81 /*
82  * Helper function for rotations:
83  * - old's parent and color get assigned to new
84  * - old gets assigned new as a parent and 'color' as a color.
85  */
86 static inline void
87 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
88                         struct rb_root *root, int color)
89 {
90         struct rb_node *parent = rb_parent(old);
91         new->__rb_parent_color = old->__rb_parent_color;
92         rb_set_parent_color(old, new, color);
93         __rb_change_child(old, new, parent, root);
94 }
95
96 static __always_inline void
97 __rb_insert(struct rb_node *node, struct rb_root *root,
98             bool newleft, struct rb_node **leftmost,
99             void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
100 {
101         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
102
103         if (newleft)
104                 *leftmost = node;
105
106         while (true) {
107                 /*
108                  * Loop invariant: node is red.
109                  */
110                 if (unlikely(!parent)) {
111                         /*
112                          * The inserted node is root. Either this is the
113                          * first node, or we recursed at Case 1 below and
114                          * are no longer violating 4).
115                          */
116                         rb_set_parent_color(node, NULL, RB_BLACK);
117                         break;
118                 }
119
120                 /*
121                  * If there is a black parent, we are done.
122                  * Otherwise, take some corrective action as,
123                  * per 4), we don't want a red root or two
124                  * consecutive red nodes.
125                  */
126                 if(rb_is_black(parent))
127                         break;
128
129                 gparent = rb_red_parent(parent);
130
131                 tmp = gparent->rb_right;
132                 if (parent != tmp) {    /* parent == gparent->rb_left */
133                         if (tmp && rb_is_red(tmp)) {
134                                 /*
135                                  * Case 1 - node's uncle is red (color flips).
136                                  *
137                                  *       G            g
138                                  *      / \          / \
139                                  *     p   u  -->   P   U
140                                  *    /            /
141                                  *   n            n
142                                  *
143                                  * However, since g's parent might be red, and
144                                  * 4) does not allow this, we need to recurse
145                                  * at g.
146                                  */
147                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
148                                 rb_set_parent_color(parent, gparent, RB_BLACK);
149                                 node = gparent;
150                                 parent = rb_parent(node);
151                                 rb_set_parent_color(node, parent, RB_RED);
152                                 continue;
153                         }
154
155                         tmp = parent->rb_right;
156                         if (node == tmp) {
157                                 /*
158                                  * Case 2 - node's uncle is black and node is
159                                  * the parent's right child (left rotate at parent).
160                                  *
161                                  *      G             G
162                                  *     / \           / \
163                                  *    p   U  -->    n   U
164                                  *     \           /
165                                  *      n         p
166                                  *
167                                  * This still leaves us in violation of 4), the
168                                  * continuation into Case 3 will fix that.
169                                  */
170                                 tmp = node->rb_left;
171                                 WRITE_ONCE(parent->rb_right, tmp);
172                                 WRITE_ONCE(node->rb_left, parent);
173                                 if (tmp)
174                                         rb_set_parent_color(tmp, parent,
175                                                             RB_BLACK);
176                                 rb_set_parent_color(parent, node, RB_RED);
177                                 augment_rotate(parent, node);
178                                 parent = node;
179                                 tmp = node->rb_right;
180                         }
181
182                         /*
183                          * Case 3 - node's uncle is black and node is
184                          * the parent's left child (right rotate at gparent).
185                          *
186                          *        G           P
187                          *       / \         / \
188                          *      p   U  -->  n   g
189                          *     /                 \
190                          *    n                   U
191                          */
192                         WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
193                         WRITE_ONCE(parent->rb_right, gparent);
194                         if (tmp)
195                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
196                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
197                         augment_rotate(gparent, parent);
198                         break;
199                 } else {
200                         tmp = gparent->rb_left;
201                         if (tmp && rb_is_red(tmp)) {
202                                 /* Case 1 - color flips */
203                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
204                                 rb_set_parent_color(parent, gparent, RB_BLACK);
205                                 node = gparent;
206                                 parent = rb_parent(node);
207                                 rb_set_parent_color(node, parent, RB_RED);
208                                 continue;
209                         }
210
211                         tmp = parent->rb_left;
212                         if (node == tmp) {
213                                 /* Case 2 - right rotate at parent */
214                                 tmp = node->rb_right;
215                                 WRITE_ONCE(parent->rb_left, tmp);
216                                 WRITE_ONCE(node->rb_right, parent);
217                                 if (tmp)
218                                         rb_set_parent_color(tmp, parent,
219                                                             RB_BLACK);
220                                 rb_set_parent_color(parent, node, RB_RED);
221                                 augment_rotate(parent, node);
222                                 parent = node;
223                                 tmp = node->rb_left;
224                         }
225
226                         /* Case 3 - left rotate at gparent */
227                         WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
228                         WRITE_ONCE(parent->rb_left, gparent);
229                         if (tmp)
230                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
231                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
232                         augment_rotate(gparent, parent);
233                         break;
234                 }
235         }
236 }
237
238 /*
239  * Inline version for rb_erase() use - we want to be able to inline
240  * and eliminate the dummy_rotate callback there
241  */
242 static __always_inline void
243 ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
244         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
245 {
246         struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
247
248         while (true) {
249                 /*
250                  * Loop invariants:
251                  * - node is black (or NULL on first iteration)
252                  * - node is not the root (parent is not NULL)
253                  * - All leaf paths going through parent and node have a
254                  *   black node count that is 1 lower than other leaf paths.
255                  */
256                 sibling = parent->rb_right;
257                 if (node != sibling) {  /* node == parent->rb_left */
258                         if (rb_is_red(sibling)) {
259                                 /*
260                                  * Case 1 - left rotate at parent
261                                  *
262                                  *     P               S
263                                  *    / \             / \
264                                  *   N   s    -->    p   Sr
265                                  *      / \         / \
266                                  *     Sl  Sr      N   Sl
267                                  */
268                                 tmp1 = sibling->rb_left;
269                                 WRITE_ONCE(parent->rb_right, tmp1);
270                                 WRITE_ONCE(sibling->rb_left, parent);
271                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
272                                 __rb_rotate_set_parents(parent, sibling, root,
273                                                         RB_RED);
274                                 augment_rotate(parent, sibling);
275                                 sibling = tmp1;
276                         }
277                         tmp1 = sibling->rb_right;
278                         if (!tmp1 || rb_is_black(tmp1)) {
279                                 tmp2 = sibling->rb_left;
280                                 if (!tmp2 || rb_is_black(tmp2)) {
281                                         /*
282                                          * Case 2 - sibling color flip
283                                          * (p could be either color here)
284                                          *
285                                          *    (p)           (p)
286                                          *    / \           / \
287                                          *   N   S    -->  N   s
288                                          *      / \           / \
289                                          *     Sl  Sr        Sl  Sr
290                                          *
291                                          * This leaves us violating 5) which
292                                          * can be fixed by flipping p to black
293                                          * if it was red, or by recursing at p.
294                                          * p is red when coming from Case 1.
295                                          */
296                                         rb_set_parent_color(sibling, parent,
297                                                             RB_RED);
298                                         if (rb_is_red(parent))
299                                                 rb_set_black(parent);
300                                         else {
301                                                 node = parent;
302                                                 parent = rb_parent(node);
303                                                 if (parent)
304                                                         continue;
305                                         }
306                                         break;
307                                 }
308                                 /*
309                                  * Case 3 - right rotate at sibling
310                                  * (p could be either color here)
311                                  *
312                                  *   (p)           (p)
313                                  *   / \           / \
314                                  *  N   S    -->  N   sl
315                                  *     / \             \
316                                  *    sl  Sr            S
317                                  *                       \
318                                  *                        Sr
319                                  *
320                                  * Note: p might be red, and then both
321                                  * p and sl are red after rotation(which
322                                  * breaks property 4). This is fixed in
323                                  * Case 4 (in __rb_rotate_set_parents()
324                                  *         which set sl the color of p
325                                  *         and set p RB_BLACK)
326                                  *
327                                  *   (p)            (sl)
328                                  *   / \            /  \
329                                  *  N   sl   -->   P    S
330                                  *       \        /      \
331                                  *        S      N        Sr
332                                  *         \
333                                  *          Sr
334                                  */
335                                 tmp1 = tmp2->rb_right;
336                                 WRITE_ONCE(sibling->rb_left, tmp1);
337                                 WRITE_ONCE(tmp2->rb_right, sibling);
338                                 WRITE_ONCE(parent->rb_right, tmp2);
339                                 if (tmp1)
340                                         rb_set_parent_color(tmp1, sibling,
341                                                             RB_BLACK);
342                                 augment_rotate(sibling, tmp2);
343                                 tmp1 = sibling;
344                                 sibling = tmp2;
345                         }
346                         /*
347                          * Case 4 - left rotate at parent + color flips
348                          * (p and sl could be either color here.
349                          *  After rotation, p becomes black, s acquires
350                          *  p's color, and sl keeps its color)
351                          *
352                          *      (p)             (s)
353                          *      / \             / \
354                          *     N   S     -->   P   Sr
355                          *        / \         / \
356                          *      (sl) sr      N  (sl)
357                          */
358                         tmp2 = sibling->rb_left;
359                         WRITE_ONCE(parent->rb_right, tmp2);
360                         WRITE_ONCE(sibling->rb_left, parent);
361                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
362                         if (tmp2)
363                                 rb_set_parent(tmp2, parent);
364                         __rb_rotate_set_parents(parent, sibling, root,
365                                                 RB_BLACK);
366                         augment_rotate(parent, sibling);
367                         break;
368                 } else {
369                         sibling = parent->rb_left;
370                         if (rb_is_red(sibling)) {
371                                 /* Case 1 - right rotate at parent */
372                                 tmp1 = sibling->rb_right;
373                                 WRITE_ONCE(parent->rb_left, tmp1);
374                                 WRITE_ONCE(sibling->rb_right, parent);
375                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
376                                 __rb_rotate_set_parents(parent, sibling, root,
377                                                         RB_RED);
378                                 augment_rotate(parent, sibling);
379                                 sibling = tmp1;
380                         }
381                         tmp1 = sibling->rb_left;
382                         if (!tmp1 || rb_is_black(tmp1)) {
383                                 tmp2 = sibling->rb_right;
384                                 if (!tmp2 || rb_is_black(tmp2)) {
385                                         /* Case 2 - sibling color flip */
386                                         rb_set_parent_color(sibling, parent,
387                                                             RB_RED);
388                                         if (rb_is_red(parent))
389                                                 rb_set_black(parent);
390                                         else {
391                                                 node = parent;
392                                                 parent = rb_parent(node);
393                                                 if (parent)
394                                                         continue;
395                                         }
396                                         break;
397                                 }
398                                 /* Case 3 - left rotate at sibling */
399                                 tmp1 = tmp2->rb_left;
400                                 WRITE_ONCE(sibling->rb_right, tmp1);
401                                 WRITE_ONCE(tmp2->rb_left, sibling);
402                                 WRITE_ONCE(parent->rb_left, tmp2);
403                                 if (tmp1)
404                                         rb_set_parent_color(tmp1, sibling,
405                                                             RB_BLACK);
406                                 augment_rotate(sibling, tmp2);
407                                 tmp1 = sibling;
408                                 sibling = tmp2;
409                         }
410                         /* Case 4 - right rotate at parent + color flips */
411                         tmp2 = sibling->rb_right;
412                         WRITE_ONCE(parent->rb_left, tmp2);
413                         WRITE_ONCE(sibling->rb_right, parent);
414                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
415                         if (tmp2)
416                                 rb_set_parent(tmp2, parent);
417                         __rb_rotate_set_parents(parent, sibling, root,
418                                                 RB_BLACK);
419                         augment_rotate(parent, sibling);
420                         break;
421                 }
422         }
423 }
424
425 /* Non-inline version for rb_erase_augmented() use */
426 void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
427         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
428 {
429         ____rb_erase_color(parent, root, augment_rotate);
430 }
431 EXPORT_SYMBOL(__rb_erase_color);
432
433 /*
434  * Non-augmented rbtree manipulation functions.
435  *
436  * We use dummy augmented callbacks here, and have the compiler optimize them
437  * out of the rb_insert_color() and rb_erase() function definitions.
438  */
439
440 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
441 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
442 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
443
444 static const struct rb_augment_callbacks dummy_callbacks = {
445         .propagate = dummy_propagate,
446         .copy = dummy_copy,
447         .rotate = dummy_rotate
448 };
449
450 void rb_insert_color(struct rb_node *node, struct rb_root *root)
451 {
452         __rb_insert(node, root, false, NULL, dummy_rotate);
453 }
454 EXPORT_SYMBOL(rb_insert_color);
455
456 void rb_erase(struct rb_node *node, struct rb_root *root)
457 {
458         struct rb_node *rebalance;
459         rebalance = __rb_erase_augmented(node, root,
460                                          NULL, &dummy_callbacks);
461         if (rebalance)
462                 ____rb_erase_color(rebalance, root, dummy_rotate);
463 }
464 EXPORT_SYMBOL(rb_erase);
465
466 void rb_insert_color_cached(struct rb_node *node,
467                             struct rb_root_cached *root, bool leftmost)
468 {
469         __rb_insert(node, &root->rb_root, leftmost,
470                     &root->rb_leftmost, dummy_rotate);
471 }
472 EXPORT_SYMBOL(rb_insert_color_cached);
473
474 void rb_erase_cached(struct rb_node *node, struct rb_root_cached *root)
475 {
476         struct rb_node *rebalance;
477         rebalance = __rb_erase_augmented(node, &root->rb_root,
478                                          &root->rb_leftmost, &dummy_callbacks);
479         if (rebalance)
480                 ____rb_erase_color(rebalance, &root->rb_root, dummy_rotate);
481 }
482 EXPORT_SYMBOL(rb_erase_cached);
483
484 /*
485  * Augmented rbtree manipulation functions.
486  *
487  * This instantiates the same __always_inline functions as in the non-augmented
488  * case, but this time with user-defined callbacks.
489  */
490
491 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
492                            bool newleft, struct rb_node **leftmost,
493         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
494 {
495         __rb_insert(node, root, newleft, leftmost, augment_rotate);
496 }
497 EXPORT_SYMBOL(__rb_insert_augmented);
498
499 /*
500  * This function returns the first node (in sort order) of the tree.
501  */
502 struct rb_node *rb_first(const struct rb_root *root)
503 {
504         struct rb_node  *n;
505
506         n = root->rb_node;
507         if (!n)
508                 return NULL;
509         while (n->rb_left)
510                 n = n->rb_left;
511         return n;
512 }
513 EXPORT_SYMBOL(rb_first);
514
515 struct rb_node *rb_last(const struct rb_root *root)
516 {
517         struct rb_node  *n;
518
519         n = root->rb_node;
520         if (!n)
521                 return NULL;
522         while (n->rb_right)
523                 n = n->rb_right;
524         return n;
525 }
526 EXPORT_SYMBOL(rb_last);
527
528 struct rb_node *rb_next(const struct rb_node *node)
529 {
530         struct rb_node *parent;
531
532         if (RB_EMPTY_NODE(node))
533                 return NULL;
534
535         /*
536          * If we have a right-hand child, go down and then left as far
537          * as we can.
538          */
539         if (node->rb_right) {
540                 node = node->rb_right;
541                 while (node->rb_left)
542                         node=node->rb_left;
543                 return (struct rb_node *)node;
544         }
545
546         /*
547          * No right-hand children. Everything down and left is smaller than us,
548          * so any 'next' node must be in the general direction of our parent.
549          * Go up the tree; any time the ancestor is a right-hand child of its
550          * parent, keep going up. First time it's a left-hand child of its
551          * parent, said parent is our 'next' node.
552          */
553         while ((parent = rb_parent(node)) && node == parent->rb_right)
554                 node = parent;
555
556         return parent;
557 }
558 EXPORT_SYMBOL(rb_next);
559
560 struct rb_node *rb_prev(const struct rb_node *node)
561 {
562         struct rb_node *parent;
563
564         if (RB_EMPTY_NODE(node))
565                 return NULL;
566
567         /*
568          * If we have a left-hand child, go down and then right as far
569          * as we can.
570          */
571         if (node->rb_left) {
572                 node = node->rb_left;
573                 while (node->rb_right)
574                         node=node->rb_right;
575                 return (struct rb_node *)node;
576         }
577
578         /*
579          * No left-hand children. Go up till we find an ancestor which
580          * is a right-hand child of its parent.
581          */
582         while ((parent = rb_parent(node)) && node == parent->rb_left)
583                 node = parent;
584
585         return parent;
586 }
587 EXPORT_SYMBOL(rb_prev);
588
589 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
590                      struct rb_root *root)
591 {
592         struct rb_node *parent = rb_parent(victim);
593
594         /* Copy the pointers/colour from the victim to the replacement */
595         *new = *victim;
596
597         /* Set the surrounding nodes to point to the replacement */
598         if (victim->rb_left)
599                 rb_set_parent(victim->rb_left, new);
600         if (victim->rb_right)
601                 rb_set_parent(victim->rb_right, new);
602         __rb_change_child(victim, new, parent, root);
603 }
604 EXPORT_SYMBOL(rb_replace_node);
605
606 void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new,
607                          struct rb_root *root)
608 {
609         struct rb_node *parent = rb_parent(victim);
610
611         /* Copy the pointers/colour from the victim to the replacement */
612         *new = *victim;
613
614         /* Set the surrounding nodes to point to the replacement */
615         if (victim->rb_left)
616                 rb_set_parent(victim->rb_left, new);
617         if (victim->rb_right)
618                 rb_set_parent(victim->rb_right, new);
619
620         /* Set the parent's pointer to the new node last after an RCU barrier
621          * so that the pointers onwards are seen to be set correctly when doing
622          * an RCU walk over the tree.
623          */
624         __rb_change_child_rcu(victim, new, parent, root);
625 }
626 EXPORT_SYMBOL(rb_replace_node_rcu);
627
628 static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
629 {
630         for (;;) {
631                 if (node->rb_left)
632                         node = node->rb_left;
633                 else if (node->rb_right)
634                         node = node->rb_right;
635                 else
636                         return (struct rb_node *)node;
637         }
638 }
639
640 struct rb_node *rb_next_postorder(const struct rb_node *node)
641 {
642         const struct rb_node *parent;
643         if (!node)
644                 return NULL;
645         parent = rb_parent(node);
646
647         /* If we're sitting on node, we've already seen our children */
648         if (parent && node == parent->rb_left && parent->rb_right) {
649                 /* If we are the parent's left node, go to the parent's right
650                  * node then all the way down to the left */
651                 return rb_left_deepest_node(parent->rb_right);
652         } else
653                 /* Otherwise we are the parent's right node, and the parent
654                  * should be next */
655                 return (struct rb_node *)parent;
656 }
657 EXPORT_SYMBOL(rb_next_postorder);
658
659 struct rb_node *rb_first_postorder(const struct rb_root *root)
660 {
661         if (!root->rb_node)
662                 return NULL;
663
664         return rb_left_deepest_node(root->rb_node);
665 }
666 EXPORT_SYMBOL(rb_first_postorder);