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 }