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 rotateright(struct opr_rbtree *head, struct opr_rbtree_node *node)
 164 {
 165     struct opr_rbtree_node *left = node->left;
 166
 167     node->left = left->right;
 168     if (left->right)
 169         left->right->parent = node;
 170
 171     left->right = node;
 172     left->parent = node->parent;
 173
 174     update_parent_ptr(head, node, left);
 175
 176     node->parent = left;
 177 }
 178
 179 static_inline void
 180 rotateleft(struct opr_rbtree *head, struct opr_rbtree_node *node)
 181 {
 182     struct opr_rbtree_node *right = node->right;
 183
 184     node->right = right->left;
 185     if (right->left)
 186         right->left->parent = node;
 187
 188     right->left = node;
 189     right->parent = node->parent;
 190
 191     update_parent_ptr(head, node, right);
 192
 193     node->parent = right;
 194 }
 195
 196 static_inline void
 197 swapnode(struct opr_rbtree_node **a, struct opr_rbtree_node **b)
 198 {
 199     struct opr_rbtree_node *tmp;
 200
 201     tmp = *a;
 202     *a = *b;
 203     *b = tmp;
 204 }
 205
 206 static void
 207 insert_recolour(struct opr_rbtree *head, struct opr_rbtree_node *node)
 208 {
 209     struct opr_rbtree_node *parent, *gramps;
 210
 211     while ((parent = node->parent) && parent->red) {
 212         gramps = parent->parent;
 213
 214         if (parent == gramps->left) {
 215             struct opr_rbtree_node *uncle = gramps->right;
 216
 217             if (uncle && uncle->red) {
 218                 uncle->red = 0;
 219                 parent->red = 0;
 220                 gramps->red = 1;
 221                 node = gramps;
 222                 continue;
 223             }
 224
 225             if (parent->right == node) {
 226                 rotateleft(head, parent);
 227                 swapnode(&parent, &node);
 228             }
 229
 230             parent->red = 0;
 231             gramps->red = 1;
 232             rotateright(head, gramps);
 233         } else {
 234             struct opr_rbtree_node *uncle = gramps->left;
 235
 236             if (uncle && uncle->red) {
 237                 uncle->red = 0;
 238                 parent->red = 0;
 239                 gramps->red = 1;
 240                 node = gramps;
 241                 continue;
 242             }
 243
 244             if (parent->left == node) {
 245                 rotateright(head, parent);
 246                 swapnode(&parent, &node);
 247             }
 248
 249             parent->red = 0;
 250             gramps->red = 1;
 251             rotateleft(head, gramps);
 252         }
 253     }
 254
 255     head->root->red = 0;
 256 }
 257
 258 void
 259 opr_rbtree_insert(struct opr_rbtree *head,
 260                   struct opr_rbtree_node *parent,
 261                   struct opr_rbtree_node **childptr,
 262                   struct opr_rbtree_node *node)
 263 {
 264     /* Link node 'node' into the tree at position 'parent', using either the
 265      * left or right pointers */
 266
 267     node->parent = parent;
 268     node->left = node->right = NULL;
 269     node->red = 1;
 270     *childptr = node;
 271
 272     /* Rebalance the tree for the newly inserted node */
 273     insert_recolour(head, node);
 274 }
 275
 276 static void
 277 remove_recolour(struct opr_rbtree *head, struct opr_rbtree_node *parent,
 278                 struct opr_rbtree_node *node)
 279 {
 280     struct opr_rbtree_node *other;
 281
 282     while ((node == NULL || !node->red) && node != head->root) {
 283         if (parent->left == node) {
 284             other = parent->right;
 285             if (other->red) {
 286                 other->red = 0;
 287                 parent->red = 1;
 288                 rotateleft(head, parent);
 289                 other = parent->right;
 290             }
 291             if ((other->left == NULL || !other->left->red) &&
 292                 (other->right == NULL || !other->right->red)) {
 293                 other->red = 1;
 294                 node = parent;
 295                 parent = node->parent;
 296             } else {
 297                 if (other->right == NULL || !other->right->red) {
 298                     other->left->red = 0;
 299                     other->red = 1;
 300                     rotateright(head, other);
 301                     other = parent->right;
 302                 }
 303                 other->red = parent->red;
 304                 parent->red = 0;
 305                 other->right->red = 0;
 306                 rotateleft(head, parent);
 307                 node = head->root;
 308                 break;
 309             }
 310         } else {
 311             other = parent->left;
 312             if (other->red) {
 313                 other->red = 0;
 314                 parent->red = 1;
 315                 rotateright(head, parent);
 316                 other = parent->left;
 317             }
 318             if ((other->left == NULL || !other->left->red) &&
 319                 (other->right == NULL || !other->right->red)) {
 320                 other->red = 1;
 321                 node = parent;
 322                 parent = node->parent;
 323             } else {
 324                 if (other->left == NULL || !other->left->red) {
 325                     other->right->red = 0;
 326                     other->red = 1;
 327                     rotateleft(head, other);
 328                     other = parent->left;
 329                 }
 330                 other->red = parent->red;
 331                 parent->red = 0;
 332                 other->left->red = 0;
 333                 rotateright(head, parent);
 334                 node = head->root;
 335                 break;
 336             }
 337         }
 338     }
 339     if (node)
 340         node->red = 0;
 341 }
 342
 343 void
 344 opr_rbtree_remove(struct opr_rbtree *head, struct opr_rbtree_node *node)
 345 {
 346     struct opr_rbtree_node *child, *parent;
 347     int red;
 348
 349
 350     if (node->left == NULL && node->right == NULL) {
 351         /* A node with no non-leaf children */
 352         update_parent_ptr(head, node, NULL);
 353
 354         if (!node->red)
 355             remove_recolour(head, node->parent, NULL);
 356
 357         return;
 358     }
 359
 360     if (node->left != NULL && node->right != NULL) {
 361         /* A node with two children.
 362          *
 363          * Move the next node in the tree (which will be a leaf node)
 364          * onto our tree current position, then rebalance as required
 365          */
 366         struct opr_rbtree_node *old, *left;
 367
 368         old = node;
 369
 370         /* Set node to the next node in the tree from the current
 371          * position, where the next node is the left-most leaf node
 372          * in our right child */
 373         node = node->right;
 374         while ((left = node->left) != NULL)
 375             node = left;
 376
 377         /* Move 'node' into the position occupied by 'old', which is being
 378          * removed */
 379
 380         update_parent_ptr(head, old, node);
 381
 382         child = node->right;
 383         parent = node->parent;
 384         red = node->red;
 385
 386         /* As we're logically just copying the value, must preserve the
 387          * old node's colour */
 388         node->red = old->red;
 389
 390         /* ... and the old node's linkage */
 391         if (parent == old)
 392             parent = node;
 393         else {
 394             if (child)
 395                 child->parent = parent;
 396             parent->left = child;
 397
 398             node->right = old->right;
 399             old->right->parent = node;
 400         }
 401
 402         node->parent = old->parent;
 403         node->left = old->left;
 404         old->left->parent = node;
 405
 406         /* If the node being removed was black, then we must recolour the
 407          * tree to maintain balance */
 408         if (!red)
 409             remove_recolour(head, parent, child);
 410
 411         return;
 412     }
 413
 414     /* Only remaining option - node with a single child */
 415
 416     if (node->left == NULL)
 417         child = node->right;
 418     else
 419         child = node->left;
 420
 421     child->parent = node->parent;
 422
 423     update_parent_ptr(head, node, child);
 424
 425     if (!node->red)
 426         remove_recolour(head, node->parent, child);
 427 }
 428
 429 void
 430 opr_rbtree_replace(struct opr_rbtree *head,
 431                    struct opr_rbtree_node *old,
 432                    struct opr_rbtree_node *replacement)
 433 {
 434     /* Update our parent's pointer to us */
 435     update_parent_ptr(head, old, replacement);
 436
 437     /* And our children's pointers to us */
 438     if (old->left)
 439         old->left->parent = replacement;
 440     if (old->right)
 441         old->right->parent = replacement;
 442
 443     /* Copy over parent, left, right and colour */
 444     *replacement = *old;
 445 }