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