unbound/util/rbtree.c

   1 /*
   2  * rbtree.c -- generic red black tree
   3  *
   4  * Copyright (c) 2001-2007, NLnet Labs. All rights reserved.
   5  *
   6  * This software is open source.
   7  *
   8  * Redistribution and use in source and binary forms, with or without
   9  * modification, are permitted provided that the following conditions
  10  * are met:
  11  *
  12  * Redistributions of source code must retain the above copyright notice,
  13  * this list of conditions and the following disclaimer.
  14  *
  15  * Redistributions in binary form must reproduce the above copyright notice,
  16  * this list of conditions and the following disclaimer in the documentation
  17  * and/or other materials provided with the distribution.
  18  *
  19  * Neither the name of the NLNET LABS nor the names of its contributors may
  20  * be used to endorse or promote products derived from this software without
  21  * specific prior written permission.
  22  *
  23  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  24  * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  25  * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  26  * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  27  * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  28  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
  29  * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  30  * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  31  * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  32  * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  33  * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  34  *
  35  */
  36
  37 /**
  38  * \file
  39  * Implementation of a redblack tree.
  40  */
  41
  42 #include "config.h"
  43 #include "log.h"
  44 #include "fptr_wlist.h"
  45 #include "util/rbtree.h"
  46
  47 /** Node colour black */
  48 #define BLACK   0
  49 /** Node colour red */
  50 #define RED     1
  51
  52 /** the NULL node, global alloc */
  53 rbnode_t        rbtree_null_node = {
  54         RBTREE_NULL,            /* Parent.  */
  55         RBTREE_NULL,            /* Left.  */
  56         RBTREE_NULL,            /* Right.  */
  57         NULL,                   /* Key.  */
  58         BLACK                   /* Color.  */
  59 };
  60
  61 /** rotate subtree left (to preserve redblack property) */
  62 static void rbtree_rotate_left(rbtree_t *rbtree, rbnode_t *node);
  63 /** rotate subtree right (to preserve redblack property) */
  64 static void rbtree_rotate_right(rbtree_t *rbtree, rbnode_t *node);
  65 /** Fixup node colours when insert happened */
  66 static void rbtree_insert_fixup(rbtree_t *rbtree, rbnode_t *node);
  67 /** Fixup node colours when delete happened */
  68 static void rbtree_delete_fixup(rbtree_t* rbtree, rbnode_t* child, rbnode_t* child_parent);
  69
  70 /*
  71  * Creates a new red black tree, intializes and returns a pointer to it.
  72  *
  73  * Return NULL on failure.
  74  *
  75  */
  76 rbtree_t *
  77 rbtree_create (int (*cmpf)(const void *, const void *))
  78 {
  79         rbtree_t *rbtree;
  80
  81         /* Allocate memory for it */
  82         rbtree = (rbtree_t *) malloc(sizeof(rbtree_t));
  83         if (!rbtree) {
  84                 return NULL;
  85         }
  86
  87         /* Initialize it */
  88         rbtree_init(rbtree, cmpf);
  89
  90         return rbtree;
  91 }
  92
  93 void
  94 rbtree_init(rbtree_t *rbtree, int (*cmpf)(const void *, const void *))
  95 {
  96         /* Initialize it */
  97         rbtree->root = RBTREE_NULL;
  98         rbtree->count = 0;
  99         rbtree->cmp = cmpf;
 100 }
 101
 102 /*
 103  * Rotates the node to the left.
 104  *
 105  */
 106 static void
 107 rbtree_rotate_left(rbtree_t *rbtree, rbnode_t *node)
 108 {
 109         rbnode_t *right = node->right;
 110         node->right = right->left;
 111         if (right->left != RBTREE_NULL)
 112                 right->left->parent = node;
 113
 114         right->parent = node->parent;
 115
 116         if (node->parent != RBTREE_NULL) {
 117                 if (node == node->parent->left) {
 118                         node->parent->left = right;
 119                 } else  {
 120                         node->parent->right = right;
 121                 }
 122         } else {
 123                 rbtree->root = right;
 124         }
 125         right->left = node;
 126         node->parent = right;
 127 }
 128
 129 /*
 130  * Rotates the node to the right.
 131  *
 132  */
 133 static void
 134 rbtree_rotate_right(rbtree_t *rbtree, rbnode_t *node)
 135 {
 136         rbnode_t *left = node->left;
 137         node->left = left->right;
 138         if (left->right != RBTREE_NULL)
 139                 left->right->parent = node;
 140
 141         left->parent = node->parent;
 142
 143         if (node->parent != RBTREE_NULL) {
 144                 if (node == node->parent->right) {
 145                         node->parent->right = left;
 146                 } else  {
 147                         node->parent->left = left;
 148                 }
 149         } else {
 150                 rbtree->root = left;
 151         }
 152         left->right = node;
 153         node->parent = left;
 154 }
 155
 156 static void
 157 rbtree_insert_fixup(rbtree_t *rbtree, rbnode_t *node)
 158 {
 159         rbnode_t        *uncle;
 160
 161         /* While not at the root and need fixing... */
 162         while (node != rbtree->root && node->parent->color == RED) {
 163                 /* If our parent is left child of our grandparent... */
 164                 if (node->parent == node->parent->parent->left) {
 165                         uncle = node->parent->parent->right;
 166
 167                         /* If our uncle is red... */
 168                         if (uncle->color == RED) {
 169                                 /* Paint the parent and the uncle black... */
 170                                 node->parent->color = BLACK;
 171                                 uncle->color = BLACK;
 172
 173                                 /* And the grandparent red... */
 174                                 node->parent->parent->color = RED;
 175
 176                                 /* And continue fixing the grandparent */
 177                                 node = node->parent->parent;
 178                         } else {                                /* Our uncle is black... */
 179                                 /* Are we the right child? */
 180                                 if (node == node->parent->right) {
 181                                         node = node->parent;
 182                                         rbtree_rotate_left(rbtree, node);
 183                                 }
 184                                 /* Now we're the left child, repaint and rotate... */
 185                                 node->parent->color = BLACK;
 186                                 node->parent->parent->color = RED;
 187                                 rbtree_rotate_right(rbtree, node->parent->parent);
 188                         }
 189                 } else {
 190                         uncle = node->parent->parent->left;
 191
 192                         /* If our uncle is red... */
 193                         if (uncle->color == RED) {
 194                                 /* Paint the parent and the uncle black... */
 195                                 node->parent->color = BLACK;
 196                                 uncle->color = BLACK;
 197
 198                                 /* And the grandparent red... */
 199                                 node->parent->parent->color = RED;
 200
 201                                 /* And continue fixing the grandparent */
 202                                 node = node->parent->parent;
 203                         } else {                                /* Our uncle is black... */
 204                                 /* Are we the right child? */
 205                                 if (node == node->parent->left) {
 206                                         node = node->parent;
 207                                         rbtree_rotate_right(rbtree, node);
 208                                 }
 209                                 /* Now we're the right child, repaint and rotate... */
 210                                 node->parent->color = BLACK;
 211                                 node->parent->parent->color = RED;
 212                                 rbtree_rotate_left(rbtree, node->parent->parent);
 213                         }
 214                 }
 215         }
 216         rbtree->root->color = BLACK;
 217 }
 218
 219
 220 /*
 221  * Inserts a node into a red black tree.
 222  *
 223  * Returns NULL on failure or the pointer to the newly added node
 224  * otherwise.
 225  */
 226 rbnode_t *
 227 rbtree_insert (rbtree_t *rbtree, rbnode_t *data)
 228 {
 229         /* XXX Not necessary, but keeps compiler quiet... */
 230         int r = 0;
 231
 232         /* We start at the root of the tree */
 233         rbnode_t        *node = rbtree->root;
 234         rbnode_t        *parent = RBTREE_NULL;
 235
 236         fptr_ok(fptr_whitelist_rbtree_cmp(rbtree->cmp));
 237         /* Lets find the new parent... */
 238         while (node != RBTREE_NULL) {
 239                 /* Compare two keys, do we have a duplicate? */
 240                 if ((r = rbtree->cmp(data->key, node->key)) == 0) {
 241                         return NULL;
 242                 }
 243                 parent = node;
 244
 245                 if (r < 0) {
 246                         node = node->left;
 247                 } else {
 248                         node = node->right;
 249                 }
 250         }
 251
 252         /* Initialize the new node */
 253         data->parent = parent;
 254         data->left = data->right = RBTREE_NULL;
 255         data->color = RED;
 256         rbtree->count++;
 257
 258         /* Insert it into the tree... */
 259         if (parent != RBTREE_NULL) {
 260                 if (r < 0) {
 261                         parent->left = data;
 262                 } else {
 263                         parent->right = data;
 264                 }
 265         } else {
 266                 rbtree->root = data;
 267         }
 268
 269         /* Fix up the red-black properties... */
 270         rbtree_insert_fixup(rbtree, data);
 271
 272         return data;
 273 }
 274
 275 /*
 276  * Searches the red black tree, returns the data if key is found or NULL otherwise.
 277  *
 278  */
 279 rbnode_t *
 280 rbtree_search (rbtree_t *rbtree, const void *key)
 281 {
 282         rbnode_t *node;
 283
 284         if (rbtree_find_less_equal(rbtree, key, &node)) {
 285                 return node;
 286         } else {
 287                 return NULL;
 288         }
 289 }
 290
 291 /** helpers for delete: swap node colours */
 292 static void swap_int8(uint8_t* x, uint8_t* y)
 293 {
 294         uint8_t t = *x; *x = *y; *y = t;
 295 }
 296
 297 /** helpers for delete: swap node pointers */
 298 static void swap_np(rbnode_t** x, rbnode_t** y)
 299 {
 300         rbnode_t* t = *x; *x = *y; *y = t;
 301 }
 302
 303 /** Update parent pointers of child trees of 'parent' */
 304 static void change_parent_ptr(rbtree_t* rbtree, rbnode_t* parent, rbnode_t* old, rbnode_t* new)
 305 {
 306         if(parent == RBTREE_NULL)
 307         {
 308                 log_assert(rbtree->root == old);
 309                 if(rbtree->root == old) rbtree->root = new;
 310                 return;
 311         }
 312         log_assert(parent->left == old || parent->right == old
 313                 || parent->left == new || parent->right == new);
 314         if(parent->left == old) parent->left = new;
 315         if(parent->right == old) parent->right = new;
 316 }
 317 /** Update parent pointer of a node 'child' */
 318 static void change_child_ptr(rbnode_t* child, rbnode_t* old, rbnode_t* new)
 319 {
 320         if(child == RBTREE_NULL) return;
 321         log_assert(child->parent == old || child->parent == new);
 322         if(child->parent == old) child->parent = new;
 323 }
 324
 325 rbnode_t*
 326 rbtree_delete(rbtree_t *rbtree, const void *key)
 327 {
 328         rbnode_t *to_delete;
 329         rbnode_t *child;
 330         if((to_delete = rbtree_search(rbtree, key)) == 0) return 0;
 331         rbtree->count--;
 332
 333         /* make sure we have at most one non-leaf child */
 334         if(to_delete->left != RBTREE_NULL && to_delete->right != RBTREE_NULL)
 335         {
 336                 /* swap with smallest from right subtree (or largest from left) */
 337                 rbnode_t *smright = to_delete->right;
 338                 while(smright->left != RBTREE_NULL)
 339                         smright = smright->left;
 340                 /* swap the smright and to_delete elements in the tree,
 341                  * but the rbnode_t is first part of user data struct
 342                  * so cannot just swap the keys and data pointers. Instead
 343                  * readjust the pointers left,right,parent */
 344
 345                 /* swap colors - colors are tied to the position in the tree */
 346                 swap_int8(&to_delete->color, &smright->color);
 347
 348                 /* swap child pointers in parents of smright/to_delete */
 349                 change_parent_ptr(rbtree, to_delete->parent, to_delete, smright);
 350                 if(to_delete->right != smright)
 351                         change_parent_ptr(rbtree, smright->parent, smright, to_delete);
 352
 353                 /* swap parent pointers in children of smright/to_delete */
 354                 change_child_ptr(smright->left, smright, to_delete);
 355                 change_child_ptr(smright->left, smright, to_delete);
 356                 change_child_ptr(smright->right, smright, to_delete);
 357                 change_child_ptr(smright->right, smright, to_delete);
 358                 change_child_ptr(to_delete->left, to_delete, smright);
 359                 if(to_delete->right != smright)
 360                         change_child_ptr(to_delete->right, to_delete, smright);
 361                 if(to_delete->right == smright)
 362                 {
 363                         /* set up so after swap they work */
 364                         to_delete->right = to_delete;
 365                         smright->parent = smright;
 366                 }
 367
 368                 /* swap pointers in to_delete/smright nodes */
 369                 swap_np(&to_delete->parent, &smright->parent);
 370                 swap_np(&to_delete->left, &smright->left);
 371                 swap_np(&to_delete->right, &smright->right);
 372
 373                 /* now delete to_delete (which is at the location where the smright previously was) */
 374         }
 375         log_assert(to_delete->left == RBTREE_NULL || to_delete->right == RBTREE_NULL);
 376
 377         if(to_delete->left != RBTREE_NULL) child = to_delete->left;
 378         else child = to_delete->right;
 379
 380         /* unlink to_delete from the tree, replace to_delete with child */
 381         change_parent_ptr(rbtree, to_delete->parent, to_delete, child);
 382         change_child_ptr(child, to_delete, to_delete->parent);
 383
 384         if(to_delete->color == RED)
 385         {
 386                 /* if node is red then the child (black) can be swapped in */
 387         }
 388         else if(child->color == RED)
 389         {
 390                 /* change child to BLACK, removing a RED node is no problem */
 391                 if(child!=RBTREE_NULL) child->color = BLACK;
 392         }
 393         else rbtree_delete_fixup(rbtree, child, to_delete->parent);
 394
 395         /* unlink completely */
 396         to_delete->parent = RBTREE_NULL;
 397         to_delete->left = RBTREE_NULL;
 398         to_delete->right = RBTREE_NULL;
 399         to_delete->color = BLACK;
 400         return to_delete;
 401 }
 402
 403 static void rbtree_delete_fixup(rbtree_t* rbtree, rbnode_t* child, rbnode_t* child_parent)
 404 {
 405         rbnode_t* sibling;
 406         int go_up = 1;
 407
 408         /* determine sibling to the node that is one-black short */
 409         if(child_parent->right == child) sibling = child_parent->left;
 410         else sibling = child_parent->right;
 411
 412         while(go_up)
 413         {
 414                 if(child_parent == RBTREE_NULL)
 415                 {
 416                         /* removed parent==black from root, every path, so ok */
 417                         return;
 418                 }
 419
 420                 if(sibling->color == RED)
 421                 {       /* rotate to get a black sibling */
 422                         child_parent->color = RED;
 423                         sibling->color = BLACK;
 424                         if(child_parent->right == child)
 425                                 rbtree_rotate_right(rbtree, child_parent);
 426                         else    rbtree_rotate_left(rbtree, child_parent);
 427                         /* new sibling after rotation */
 428                         if(child_parent->right == child) sibling = child_parent->left;
 429                         else sibling = child_parent->right;
 430                 }
 431
 432                 if(child_parent->color == BLACK
 433                         && sibling->color == BLACK
 434                         && sibling->left->color == BLACK
 435                         && sibling->right->color == BLACK)
 436                 {       /* fixup local with recolor of sibling */
 437                         if(sibling != RBTREE_NULL)
 438                                 sibling->color = RED;
 439
 440                         child = child_parent;
 441                         child_parent = child_parent->parent;
 442                         /* prepare to go up, new sibling */
 443                         if(child_parent->right == child) sibling = child_parent->left;
 444                         else sibling = child_parent->right;
 445                 }
 446                 else go_up = 0;
 447         }
 448
 449         if(child_parent->color == RED
 450                 && sibling->color == BLACK
 451                 && sibling->left->color == BLACK
 452                 && sibling->right->color == BLACK)
 453         {
 454                 /* move red to sibling to rebalance */
 455                 if(sibling != RBTREE_NULL)
 456                         sibling->color = RED;
 457                 child_parent->color = BLACK;
 458                 return;
 459         }
 460         log_assert(sibling != RBTREE_NULL);
 461
 462         /* get a new sibling, by rotating at sibling. See which child
 463            of sibling is red */
 464         if(child_parent->right == child
 465                 && sibling->color == BLACK
 466                 && sibling->right->color == RED
 467                 && sibling->left->color == BLACK)
 468         {
 469                 sibling->color = RED;
 470                 sibling->right->color = BLACK;
 471                 rbtree_rotate_left(rbtree, sibling);
 472                 /* new sibling after rotation */
 473                 if(child_parent->right == child) sibling = child_parent->left;
 474                 else sibling = child_parent->right;
 475         }
 476         else if(child_parent->left == child
 477                 && sibling->color == BLACK
 478                 && sibling->left->color == RED
 479                 && sibling->right->color == BLACK)
 480         {
 481                 sibling->color = RED;
 482                 sibling->left->color = BLACK;
 483                 rbtree_rotate_right(rbtree, sibling);
 484                 /* new sibling after rotation */
 485                 if(child_parent->right == child) sibling = child_parent->left;
 486                 else sibling = child_parent->right;
 487         }
 488
 489         /* now we have a black sibling with a red child. rotate and exchange colors. */
 490         sibling->color = child_parent->color;
 491         child_parent->color = BLACK;
 492         if(child_parent->right == child)
 493         {
 494                 log_assert(sibling->left->color == RED);
 495                 sibling->left->color = BLACK;
 496                 rbtree_rotate_right(rbtree, child_parent);
 497         }
 498         else
 499         {
 500                 log_assert(sibling->right->color == RED);
 501                 sibling->right->color = BLACK;
 502                 rbtree_rotate_left(rbtree, child_parent);
 503         }
 504 }
 505
 506 int
 507 rbtree_find_less_equal(rbtree_t *rbtree, const void *key, rbnode_t **result)
 508 {
 509         int r;
 510         rbnode_t *node;
 511
 512         log_assert(result);
 513
 514         /* We start at root... */
 515         node = rbtree->root;
 516
 517         *result = NULL;
 518         fptr_ok(fptr_whitelist_rbtree_cmp(rbtree->cmp));
 519
 520         /* While there are children... */
 521         while (node != RBTREE_NULL) {
 522                 r = rbtree->cmp(key, node->key);
 523                 if (r == 0) {
 524                         /* Exact match */
 525                         *result = node;
 526                         return 1;
 527                 }
 528                 if (r < 0) {
 529                         node = node->left;
 530                 } else {
 531                         /* Temporary match */
 532                         *result = node;
 533                         node = node->right;
 534                 }
 535         }
 536         return 0;
 537 }
 538
 539 /*
 540  * Finds the first element in the red black tree
 541  *
 542  */
 543 rbnode_t *
 544 rbtree_first (rbtree_t *rbtree)
 545 {
 546         rbnode_t *node;
 547
 548         for (node = rbtree->root; node->left != RBTREE_NULL; node = node->left);
 549         return node;
 550 }
 551
 552 rbnode_t *
 553 rbtree_last (rbtree_t *rbtree)
 554 {
 555         rbnode_t *node;
 556
 557         for (node = rbtree->root; node->right != RBTREE_NULL; node = node->right);
 558         return node;
 559 }
 560
 561 /*
 562  * Returns the next node...
 563  *
 564  */
 565 rbnode_t *
 566 rbtree_next (rbnode_t *node)
 567 {
 568         rbnode_t *parent;
 569
 570         if (node->right != RBTREE_NULL) {
 571                 /* One right, then keep on going left... */
 572                 for (node = node->right; node->left != RBTREE_NULL; node = node->left);
 573         } else {
 574                 parent = node->parent;
 575                 while (parent != RBTREE_NULL && node == parent->right) {
 576                         node = parent;
 577                         parent = parent->parent;
 578                 }
 579                 node = parent;
 580         }
 581         return node;
 582 }
 583
 584 rbnode_t *
 585 rbtree_previous(rbnode_t *node)
 586 {
 587         rbnode_t *parent;
 588
 589         if (node->left != RBTREE_NULL) {
 590                 /* One left, then keep on going right... */
 591                 for (node = node->left; node->right != RBTREE_NULL; node = node->right);
 592         } else {
 593                 parent = node->parent;
 594                 while (parent != RBTREE_NULL && node == parent->left) {
 595                         node = parent;
 596                         parent = parent->parent;
 597                 }
 598                 node = parent;
 599         }
 600         return node;
 601 }
 602
 603 /** recursive descent traverse */
 604 static void
 605 traverse_post(void (*func)(rbnode_t*, void*), void* arg, rbnode_t* node)
 606 {
 607         if(!node || node == RBTREE_NULL)
 608                 return;
 609         /* recurse */
 610         traverse_post(func, arg, node->left);
 611         traverse_post(func, arg, node->right);
 612         /* call user func */
 613         (*func)(node, arg);
 614 }
 615
 616 void
 617 traverse_postorder(rbtree_t* tree, void (*func)(rbnode_t*, void*), void* arg)
 618 {
 619         traverse_post(func, arg, tree->root);
 620 }