opr: Add a red/black tree implementation
[openafs.git] / src / opr / rbtree.c
1 /*
2  * Copyright (c) 2008 - 2010 Jason Evans <jasone@FreeBSD.org>
3  * Copyright (c) 2011 Your File System Inc.
4  * All rights reserved.
5  *
6  * Redistribution and use in source and binary forms, with or without
7  * modification, are permitted provided that the following conditions
8  * are met:
9  * 1. Redistributions of source code must retain the above copyright
10  *    notice, this list of conditions and the following disclaimer.
11  * 2. Redistributions in binary form must reproduce the above copyright
12  *    notice, this list of conditions and the following disclaimer in the
13  *    documentation and/or other materials provided with the distribution.
14  *
15  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR `AS IS'' AND ANY EXPRESS OR
16  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
17  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
18  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
19  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
20  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
21  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
22  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
23  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
24  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
25  */
26
27 /* Left-leaning rbtree implementation. Originally derived from the FreeBSD
28  * CPP macro definitions by Jason Evans, but extensively modified to
29  *     *) Be function, rather than macro based
30  *     *) Use parent pointers, rather than calling an embeded comparison fn
31  *     *) Use 'NULL' to represent empty nodes, rather than a magic pointer
32  */
33
34 #include <afsconfig.h>
35 #include <afs/param.h>
36
37 #ifdef KERNEL
38 # include "afs/sysincludes.h"
39 # include "afsincludes.h"
40 #else
41 # include <roken.h>
42 #endif
43
44 #include "rbtree.h"
45
46 struct {
47     struct opr_rbtree_node *left;
48     struct opr_rbtree_node *right;
49     struct opr_rbtree_node *parent;
50     int red;
51 } opr_rbtree_node;
52
53 struct {
54     struct opr_rbtree_node *root;
55 } opr_rbtree;
56
57 /* Update the parent pointers for a node which is being replaced.
58  *
59  * If the original node had no parent, then it was the root of the tree,
60  * and the replacement is too.
61  * Otherwise, the original node must have been either the left or right
62  * child of its parent, so update the left (or right) pointer to point
63  * to the replacement as appropriate.
64  */
65
66 static inline void
67 update_parent_ptr(struct opr_rbtree *head, struct opr_rbtree_node *old,
68                   struct opr_rbtree_node *replacement)
69 {
70     if (old->parent) {
71         if (old->parent->left == old)
72             old->parent->left = replacement;
73         else
74             old->parent->right = replacement;
75     } else
76         head->root = replacement;
77 }
78
79 void
80 opr_rbtree_init(struct opr_rbtree *head)
81 {
82     head->root = NULL;
83 }
84
85 struct opr_rbtree_node *
86 opr_rbtree_first(struct opr_rbtree *head)
87 {
88     struct opr_rbtree_node *node;
89
90     node = head->root;
91     if (node == NULL)
92         return node;
93
94     while (node->left != NULL)
95         node = node->left;
96
97     return node;
98 }
99
100 struct opr_rbtree_node *
101 opr_rbtree_last(struct opr_rbtree *head)
102 {
103     struct opr_rbtree_node *node;
104
105     node = head->root;
106
107     if (node == NULL)
108         return node;
109
110     while (node->right != NULL)
111         node = node->right;
112
113     return node;
114 }
115
116
117 struct opr_rbtree_node *
118 opr_rbtree_next(struct opr_rbtree_node *node)
119 {
120     struct opr_rbtree_node *parent;
121
122     /* Where there is a right child, the next node is to the right, then
123      * left as far as we can go */
124     if (node->right != NULL) {
125         node = node->right;
126         while (node->left != NULL)
127              node = node->left;
128
129         return node;
130     }
131
132     /* If there is no right hand child, then the next node is above us.
133      * Whenever our ancestor is a right-hand child, the next node is
134      * further up. When it is a left-hand child, it is our next node
135      */
136     while ((parent = node->parent) != NULL && node == parent->right)
137         node = parent;
138
139     return parent;
140 }
141
142 struct opr_rbtree_node *
143 opr_rbtree_prev(struct opr_rbtree_node *node)
144 {
145     struct opr_rbtree_node *parent;
146
147     if (node->left != NULL) {
148         node = node->left;
149         while (node->right != NULL)
150             node = node->right;
151
152         return node;
153     }
154
155     /* Same ancestor logic as for 'next', but in reverse */
156     while ((parent = node->parent) != NULL && node == parent->left)
157         node = parent;
158
159     return parent;
160 }
161
162 static inline void
163 initnode(struct opr_rbtree_node *node)
164 {
165     node->left = node->right = node->parent = NULL;
166     node->red = 1;
167 }
168
169 static inline void
170 rotateright(struct opr_rbtree *head, struct opr_rbtree_node *node)
171 {
172     struct opr_rbtree_node *left = node->left;
173
174     node->left = left->right;
175     if (left->right)
176         left->right->parent = node;
177
178     left->right = node;
179     left->parent = node->parent;
180
181     update_parent_ptr(head, node, left);
182
183     node->parent = left;
184 }
185
186 static inline void
187 rotateleft(struct opr_rbtree *head, struct opr_rbtree_node *node)
188 {
189     struct opr_rbtree_node *right = node->right;
190
191     node->right = right->left;
192     if (right->left)
193         right->left->parent = node;
194
195     right->left = node;
196     right->parent = node->parent;
197
198     update_parent_ptr(head, node, right);
199
200     node->parent = right;
201 }
202
203 static inline void
204 swapnode(struct opr_rbtree_node **a, struct opr_rbtree_node **b)
205 {
206     struct opr_rbtree_node *tmp;
207
208     tmp = *a;
209     *a = *b;
210     *b = tmp;
211 }
212
213 static void
214 insert_recolour(struct opr_rbtree *head, struct opr_rbtree_node *node)
215 {
216     struct opr_rbtree_node *parent, *gramps;
217
218     while ((parent = node->parent) && parent->red) {
219         gramps = parent->parent;
220
221         if (parent == gramps->left) {
222             struct opr_rbtree_node *uncle = gramps->right;
223
224             if (uncle && uncle->red) {
225                 uncle->red = 0;
226                 parent->red = 0;
227                 gramps->red = 1;
228                 node = gramps;
229                 continue;
230             }
231
232             if (parent->right == node) {
233                 rotateleft(head, parent);
234                 swapnode(&parent, &node);
235             }
236
237             parent->red = 0;
238             gramps->red = 1;
239             rotateright(head, gramps);
240         } else {
241             struct opr_rbtree_node *uncle = gramps->left;
242
243             if (uncle && uncle->red) {
244                 uncle->red = 0;
245                 parent->red = 0;
246                 gramps->red = 1;
247                 node = gramps;
248                 continue;
249             }
250
251             if (parent->left == node) {
252                 rotateright(head, parent);
253                 swapnode(&parent, &node);
254             }
255
256             parent->red = 0;
257             gramps->red = 1;
258             rotateleft(head, gramps);
259         }
260     }
261
262     head->root->red = 0;
263 }
264
265 void
266 opr_rbtree_insert(struct opr_rbtree *head,
267                   struct opr_rbtree_node *parent,
268                   struct opr_rbtree_node **childptr,
269                   struct opr_rbtree_node *node)
270 {
271     /* Link node 'node' into the tree at position 'parent', using either the
272      * left or right pointers */
273
274     node->parent = parent;
275     node->left = node->right = NULL;
276     node->red = 1;
277     *childptr = node;
278
279     /* Rebalance the tree for the newly inserted node */
280     insert_recolour(head, node);
281 }
282
283 static void
284 remove_recolour(struct opr_rbtree *head, struct opr_rbtree_node *parent,
285                 struct opr_rbtree_node *node)
286 {
287     struct opr_rbtree_node *other;
288
289     while ((node == NULL || !node->red) && node != head->root) {
290         if (parent->left == node) {
291             other = parent->right;
292             if (other->red) {
293                 other->red = 0;
294                 parent->red = 1;
295                 rotateleft(head, parent);
296                 other = parent->right;
297             }
298             if ((other->left == NULL || !other->left->red) &&
299                 (other->right == NULL || !other->right->red)) {
300                 other->red = 1;
301                 node = parent;
302                 parent = node->parent;
303             } else {
304                 if (other->right == NULL || !other->right->red) {
305                     other->left->red = 0;
306                     other->red = 1;
307                     rotateright(head, other);
308                     other = parent->right;
309                 }
310                 other->red = parent->red;
311                 parent->red = 0;
312                 other->right->red = 0;
313                 rotateleft(head, parent);
314                 node = head->root;
315                 break;
316             }
317         } else {
318             other = parent->left;
319             if (other->red) {
320                 other->red = 0;
321                 parent->red = 1;
322                 rotateright(head, parent);
323                 other = parent->left;
324             }
325             if ((other->left == NULL || !other->left->red) &&
326                 (other->right == NULL || !other->right->red)) {
327                 other->red = 1;
328                 node = parent;
329                 parent = node->parent;
330             } else {
331                 if (other->left == NULL || !other->left->red) {
332                     other->right->red = 0;
333                     other->red = 1;
334                     rotateleft(head, other);
335                     other = parent->left;
336                 }
337                 other->red = parent->red;
338                 parent->red = 0;
339                 other->left->red = 0;
340                 rotateright(head, parent);
341                 node = head->root;
342                 break;
343             }
344         }
345     }
346     if (node)
347         node->red = 0;
348 }
349
350 void
351 opr_rbtree_remove(struct opr_rbtree *head, struct opr_rbtree_node *node)
352 {
353     struct opr_rbtree_node *child, *parent;
354     int red;
355
356
357     if (node->left == NULL && node->right == NULL) {
358         /* A node with no non-leaf children */
359         update_parent_ptr(head, node, NULL);
360
361         if (!node->red)
362             remove_recolour(head, node->parent, NULL);
363
364         return;
365     }
366
367     if (node->left != NULL && node->right != NULL) {
368         /* A node with two children.
369          *
370          * Move the next node in the tree (which will be a leaf node)
371          * onto our tree current position, then rebalance as required
372          */
373         struct opr_rbtree_node *old, *left;
374
375         old = node;
376
377         /* Set node to the next node in the tree from the current
378          * position, where the next node is the left-most leaf node
379          * in our right child */
380         node = node->right;
381         while ((left = node->left) != NULL)
382             node = left;
383
384         /* Move 'node' into the position occupied by 'old', which is being
385          * removed */
386
387         update_parent_ptr(head, old, node);
388
389         child = node->right;
390         parent = node->parent;
391         red = node->red;
392
393         /* As we're logically just copying the value, must preserve the
394          * old node's colour */
395         node->red = old->red;
396
397         /* ... and the old node's linkage */
398         if (parent == old)
399             parent = node;
400         else {
401             if (child)
402                 child->parent = parent;
403             parent->left = child;
404
405             node->right = old->right;
406             old->right->parent = node;
407         }
408
409         node->parent = old->parent;
410         node->left = old->left;
411         old->left->parent = node;
412
413         /* If the node being removed was black, then we must recolour the
414          * tree to maintain balance */
415         if (!red)
416             remove_recolour(head, parent, child);
417
418         return;
419     }
420
421     /* Only remaining option - node with a single child */
422
423     if (node->left == NULL)
424         child = node->right;
425     else if (node->right == NULL)
426         child = node->left;
427
428     child->parent = node->parent;
429
430     update_parent_ptr(head, node, child);
431
432     if (!node->red)
433         remove_recolour(head, node->parent, child);
434 }
435
436 void
437 opr_rbtree_replace(struct opr_rbtree *head,
438                    struct opr_rbtree_node *old,
439                    struct opr_rbtree_node *replacement)
440 {
441     /* Update our parent's pointer to us */
442     update_parent_ptr(head, old, replacement);
443
444     /* And our children's pointers to us */
445     if (old->left)
446         old->left->parent = replacement;
447     if (old->right)
448         old->right->parent = replacement;
449
450     /* Copy over parent, left, right and colour */
451     *replacement = *old;
452 }