src/lib/rbtree.c

   1 /*
   2  * rbtree.c     RED-BLACK balanced binary trees.
   3  *
   4  * Version:     $Id$
   5  *
   6  *   This program is free software; you can redistribute it and/or modify
   7  *   it under the terms of the GNU General Public License as published by
   8  *   the Free Software Foundation; either version 2.1 of the License, or
   9  *   (at your option) any later version.
  10  *
  11  *   This program is distributed in the hope that it will be useful,
  12  *   but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  *   GNU General Public License for more details.
  15  *
  16  *   You should have received a copy of the GNU General Public License
  17  *   along with this program; if not, write to the Free Software
  18  *   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
  19  *
  20  *  Copyright 2004,2006  The FreeRADIUS server project
  21  */
  22
  23 RCSID("$Id$")
  24
  25 #include <freeradius-devel/libradius.h>
  26
  27 #ifdef HAVE_PTHREAD_H
  28 #include <pthread.h>
  29
  30 #define PTHREAD_MUTEX_LOCK(_x) if (_x->lock) pthread_mutex_lock(&((_x)->mutex))
  31 #define PTHREAD_MUTEX_UNLOCK(_x) if (_x->lock) pthread_mutex_unlock(&((_x)->mutex))
  32 #else
  33 #define PTHREAD_MUTEX_LOCK(_x)
  34 #define PTHREAD_MUTEX_UNLOCK(_x)
  35 #endif
  36
  37 /* Red-Black tree description */
  38 typedef enum {
  39         BLACK,
  40         RED
  41 } node_colour_t;
  42
  43 struct rbnode_t {
  44     rbnode_t            *left;          //!< Left child
  45     rbnode_t            *right;         //!< Right child
  46     rbnode_t            *parent;        //!< Parent
  47     node_colour_t       colour;         //!< Node colour (BLACK, RED)
  48     void                *data;          //!< data stored in node
  49 };
  50
  51 #define NIL &sentinel      /* all leafs are sentinels */
  52 static rbnode_t sentinel = { NIL, NIL, NULL, BLACK, NULL};
  53
  54 struct rbtree_t {
  55 #ifndef NDEBUG
  56         uint32_t                magic;
  57 #endif
  58         rbnode_t                *root;
  59         int                     num_elements;
  60         rb_comparator_t         compare;
  61         rb_free_t               free;
  62         bool                    replace;
  63 #ifdef HAVE_PTHREAD_H
  64         bool                    lock;
  65         pthread_mutex_t         mutex;
  66 #endif
  67 };
  68 #define RBTREE_MAGIC (0x5ad09c42)
  69
  70 /** Walks the tree to delete all nodes Does NOT re-balance it!
  71  *
  72  */
  73 static void free_walker(rbtree_t *tree, rbnode_t *x)
  74 {
  75         (void) talloc_get_type_abort(x, rbnode_t);
  76
  77         if (x->left != NIL) free_walker(tree, x->left);
  78         if (x->right != NIL) free_walker(tree, x->right);
  79
  80         if (tree->free) tree->free(x->data);
  81         talloc_free(x);
  82 }
  83
  84 void rbtree_free(rbtree_t *tree)
  85 {
  86         if (!tree) return;
  87
  88         PTHREAD_MUTEX_LOCK(tree);
  89
  90         /*
  91          *      walk the tree, deleting the nodes...
  92          */
  93         if (tree->root != NIL) free_walker(tree, tree->root);
  94
  95 #ifndef NDEBUG
  96         tree->magic = 0;
  97 #endif
  98         tree->root = NULL;
  99
 100 #ifdef HAVE_PTHREAD_H
 101         if (tree->lock) pthread_mutex_destroy(&tree->mutex);
 102 #endif
 103
 104         talloc_free(tree);
 105 }
 106
 107 /** Create a new RED-BLACK tree
 108  *
 109  */
 110 rbtree_t *rbtree_create(rb_comparator_t compare, rb_free_t node_free, int flags)
 111 {
 112         rbtree_t *tree;
 113
 114         if (!compare) return NULL;
 115
 116         tree = talloc_zero(NULL, rbtree_t);
 117         if (!tree) return NULL;
 118
 119 #ifndef NDEBUG
 120         tree->magic = RBTREE_MAGIC;
 121 #endif
 122         tree->root = NIL;
 123         tree->compare = compare;
 124         tree->replace = (flags & RBTREE_FLAG_REPLACE) != 0 ? true : false;
 125 #ifdef HAVE_PTHREAD_H
 126         tree->lock = (flags & RBTREE_FLAG_LOCK) != 0 ? true : false;
 127         if (tree->lock) {
 128                 pthread_mutex_init(&tree->mutex, NULL);
 129         }
 130 #endif
 131         tree->free = node_free;
 132
 133         return tree;
 134 }
 135
 136 /** Rotate Node x to left
 137  *
 138  */
 139 static void rotate_left(rbtree_t *tree, rbnode_t *x)
 140 {
 141
 142         rbnode_t *y = x->right;
 143
 144         /* establish x->right link */
 145         x->right = y->left;
 146         if (y->left != NIL) y->left->parent = x;
 147
 148         /* establish y->parent link */
 149         if (y != NIL) y->parent = x->parent;
 150         if (x->parent) {
 151                 if (x == x->parent->left) {
 152                         x->parent->left = y;
 153                 } else {
 154                         x->parent->right = y;
 155                 }
 156         } else {
 157                 tree->root = y;
 158         }
 159
 160         /* link x and y */
 161         y->left = x;
 162         if (x != NIL) x->parent = y;
 163 }
 164
 165 /** Rotate Node x to right
 166  *
 167  */
 168 static void rotate_right(rbtree_t *tree, rbnode_t *x)
 169 {
 170         rbnode_t *y = x->left;
 171
 172         /* establish x->left link */
 173         x->left = y->right;
 174         if (y->right != NIL) y->right->parent = x;
 175
 176         /* establish y->parent link */
 177         if (y != NIL) y->parent = x->parent;
 178         if (x->parent) {
 179                 if (x == x->parent->right) {
 180                         x->parent->right = y;
 181                 } else {
 182                         x->parent->left = y;
 183                 }
 184         } else {
 185                 tree->root = y;
 186         }
 187
 188         /* link x and y */
 189         y->right = x;
 190         if (x != NIL) x->parent = y;
 191 }
 192
 193 /** Maintain red-black tree balance after inserting node x
 194  *
 195  */
 196 static void insert_fixup(rbtree_t *tree, rbnode_t *x)
 197 {
 198         /* check RED-BLACK properties */
 199         while ((x != tree->root) && (x->parent->colour == RED)) {
 200                 /* we have a violation */
 201                 if (x->parent == x->parent->parent->left) {
 202                         rbnode_t *y = x->parent->parent->right;
 203                         if (y->colour == RED) {
 204
 205                                 /* uncle is RED */
 206                                 x->parent->colour = BLACK;
 207                                 y->colour = BLACK;
 208                                 x->parent->parent->colour = RED;
 209                                 x = x->parent->parent;
 210                         } else {
 211
 212                                 /* uncle is BLACK */
 213                                 if (x == x->parent->right) {
 214                                         /* make x a left child */
 215                                         x = x->parent;
 216                                         rotate_left(tree, x);
 217                                 }
 218
 219                                 /* recolour and rotate */
 220                                 x->parent->colour = BLACK;
 221                                 x->parent->parent->colour = RED;
 222                                 rotate_right(tree, x->parent->parent);
 223                         }
 224                 } else {
 225
 226                         /* mirror image of above code */
 227                         rbnode_t *y = x->parent->parent->left;
 228                         if (y->colour == RED) {
 229
 230                                 /* uncle is RED */
 231                                 x->parent->colour = BLACK;
 232                                 y->colour = BLACK;
 233                                 x->parent->parent->colour = RED;
 234                                 x = x->parent->parent;
 235                         } else {
 236
 237                                 /* uncle is BLACK */
 238                                 if (x == x->parent->left) {
 239                                         x = x->parent;
 240                                         rotate_right(tree, x);
 241                                 }
 242                                 x->parent->colour = BLACK;
 243                                 x->parent->parent->colour = RED;
 244                                 rotate_left(tree, x->parent->parent);
 245                         }
 246                 }
 247         }
 248
 249         tree->root->colour = BLACK;
 250 }
 251
 252
 253 /** Insert an element into the tree
 254  *
 255  */
 256 rbnode_t *rbtree_insert_node(rbtree_t *tree, void *data)
 257 {
 258         rbnode_t *current, *parent, *x;
 259
 260         PTHREAD_MUTEX_LOCK(tree);
 261
 262         /* find where node belongs */
 263         current = tree->root;
 264         parent = NULL;
 265         while (current != NIL) {
 266                 int result;
 267
 268                 /*
 269                  *      See if two entries are identical.
 270                  */
 271                 result = tree->compare(data, current->data);
 272                 if (result == 0) {
 273                         /*
 274                          *      Don't replace the entry.
 275                          */
 276                         if (!tree->replace) {
 277                                 PTHREAD_MUTEX_UNLOCK(tree);
 278                                 return NULL;
 279                         }
 280
 281                         /*
 282                          *      Do replace the entry.
 283                          */
 284                         if (tree->free) tree->free(current->data);
 285                         current->data = data;
 286                         PTHREAD_MUTEX_UNLOCK(tree);
 287                         return current;
 288                 }
 289
 290                 parent = current;
 291                 current = (result < 0) ? current->left : current->right;
 292         }
 293
 294         /* setup new node */
 295         x = talloc_zero(tree, rbnode_t);
 296         if (!x) {
 297                 fr_strerror_printf("No memory for new rbtree node");
 298                 PTHREAD_MUTEX_UNLOCK(tree);
 299                 return NULL;
 300         }
 301
 302         x->data = data;
 303         x->parent = parent;
 304         x->left = NIL;
 305         x->right = NIL;
 306         x->colour = RED;
 307
 308         /* insert node in tree */
 309         if (parent) {
 310                 if (tree->compare(data, parent->data) <= 0) {
 311                         parent->left = x;
 312                 } else {
 313                         parent->right = x;
 314                 }
 315         } else {
 316                 tree->root = x;
 317         }
 318
 319         insert_fixup(tree, x);
 320
 321         tree->num_elements++;
 322
 323         PTHREAD_MUTEX_UNLOCK(tree);
 324         return x;
 325 }
 326
 327 bool rbtree_insert(rbtree_t *tree, void *data)
 328 {
 329         if (rbtree_insert_node(tree, data)) return true;
 330         return false;
 331 }
 332
 333 /** Maintain RED-BLACK tree balance after deleting node x
 334  *
 335  */
 336 static void delete_fixup(rbtree_t *tree, rbnode_t *x, rbnode_t *parent)
 337 {
 338
 339         while (x != tree->root && x->colour == BLACK) {
 340                 if (x == parent->left) {
 341                         rbnode_t *w = parent->right;
 342                         if (w->colour == RED) {
 343                                 w->colour = BLACK;
 344                                 parent->colour = RED; /* parent != NIL? */
 345                                 rotate_left(tree, parent);
 346                                 w = parent->right;
 347                         }
 348                         if ((w->left->colour == BLACK) && (w->right->colour == BLACK)) {
 349                                 if (w != NIL) w->colour = RED;
 350                                 x = parent;
 351                                 parent = x->parent;
 352                         } else {
 353                                 if (w->right->colour == BLACK) {
 354                                         if (w->left != NIL) w->left->colour = BLACK;
 355                                         w->colour = RED;
 356                                         rotate_right(tree, w);
 357                                         w = parent->right;
 358                                 }
 359                                 w->colour = parent->colour;
 360                                 if (parent != NIL) parent->colour = BLACK;
 361                                 if (w->right->colour != BLACK) {
 362                                         w->right->colour = BLACK;
 363                                 }
 364                                 rotate_left(tree, parent);
 365                                 x = tree->root;
 366                         }
 367                 } else {
 368                         rbnode_t *w = parent->left;
 369                         if (w->colour == RED) {
 370                                 w->colour = BLACK;
 371                                 parent->colour = RED; /* parent != NIL? */
 372                                 rotate_right(tree, parent);
 373                                 w = parent->left;
 374                         }
 375                         if ((w->right->colour == BLACK) && (w->left->colour == BLACK)) {
 376                                 if (w != NIL) w->colour = RED;
 377                                 x = parent;
 378                                 parent = x->parent;
 379                         } else {
 380                                 if (w->left->colour == BLACK) {
 381                                         if (w->right != NIL) w->right->colour = BLACK;
 382                                         w->colour = RED;
 383                                         rotate_left(tree, w);
 384                                         w = parent->left;
 385                                 }
 386                                 w->colour = parent->colour;
 387                                 if (parent != NIL) parent->colour = BLACK;
 388                                 if (w->left->colour != BLACK) {
 389                                         w->left->colour = BLACK;
 390                                 }
 391                                 rotate_right(tree, parent);
 392                                 x = tree->root;
 393                         }
 394                 }
 395         }
 396         if (x != NIL) x->colour = BLACK; /* Avoid cache-dirty on NIL */
 397 }
 398
 399 /** Delete an element (z) from the tree
 400  *
 401  */
 402 static void rbtree_delete_internal(rbtree_t *tree, rbnode_t *z, bool skiplock)
 403 {
 404         rbnode_t *x, *y;
 405         rbnode_t *parent;
 406
 407         if (!z || z == NIL) return;
 408
 409         if (!skiplock) {
 410                 PTHREAD_MUTEX_LOCK(tree);
 411         }
 412
 413         if (z->left == NIL || z->right == NIL) {
 414                 /* y has a NIL node as a child */
 415                 y = z;
 416         } else {
 417                 /* find tree successor with a NIL node as a child */
 418                 y = z->right;
 419                 while (y->left != NIL) y = y->left;
 420         }
 421
 422         /* x is y's only child */
 423         if (y->left != NIL) {
 424                 x = y->left;
 425         } else {
 426                 x = y->right;   /* may be NIL! */
 427         }
 428
 429         /* remove y from the parent chain */
 430         parent = y->parent;
 431         if (x != NIL) x->parent = parent;
 432
 433         if (parent) {
 434                 if (y == parent->left) {
 435                         parent->left = x;
 436                 } else {
 437                         parent->right = x;
 438                 }
 439         } else {
 440                 tree->root = x;
 441         }
 442
 443         if (y != z) {
 444                 if (tree->free) tree->free(z->data);
 445                 z->data = y->data;
 446                 y->data = NULL;
 447
 448                 if ((y->colour == BLACK) && parent) {
 449                         delete_fixup(tree, x, parent);
 450                 }
 451
 452                 /*
 453                  *      The user structure in y->data MAy include a
 454                  *      pointer to y.  In that case, we CANNOT delete
 455                  *      y.  Instead, we copy z (which is now in the
 456                  *      tree) to y, and fix up the parent/child
 457                  *      pointers.
 458                  */
 459                 memcpy(y, z, sizeof(*y));
 460
 461                 if (!y->parent) {
 462                         tree->root = y;
 463                 } else {
 464                         if (y->parent->left == z) y->parent->left = y;
 465                         if (y->parent->right == z) y->parent->right = y;
 466                 }
 467                 if (y->left->parent == z) y->left->parent = y;
 468                 if (y->right->parent == z) y->right->parent = y;
 469
 470                 talloc_free(z);
 471
 472         } else {
 473                 if (tree->free) tree->free(y->data);
 474
 475                 if (y->colour == BLACK)
 476                         delete_fixup(tree, x, parent);
 477
 478                 talloc_free(y);
 479         }
 480
 481         tree->num_elements--;
 482         if (!skiplock) {
 483                 PTHREAD_MUTEX_UNLOCK(tree);
 484         }
 485 }
 486 void rbtree_delete(rbtree_t *tree, rbnode_t *z) {
 487         rbtree_delete_internal(tree, z, false);
 488 }
 489
 490 /** Delete a node from the tree, based on given data, which MUST have come from rbtree_finddata().
 491  *
 492  *
 493  */
 494 bool rbtree_deletebydata(rbtree_t *tree, void const *data)
 495 {
 496         rbnode_t *node = rbtree_find(tree, data);
 497
 498         if (!node) return false;
 499
 500         rbtree_delete(tree, node);
 501
 502         return true;
 503 }
 504
 505
 506 /** Find an element in the tree, returning the data, not the node
 507  *
 508  */
 509 rbnode_t *rbtree_find(rbtree_t *tree, void const *data)
 510 {
 511         rbnode_t *current;
 512
 513         PTHREAD_MUTEX_LOCK(tree);
 514         current = tree->root;
 515
 516         while (current != NIL) {
 517                 int result = tree->compare(data, current->data);
 518
 519                 if (result == 0) {
 520                         PTHREAD_MUTEX_UNLOCK(tree);
 521                         return current;
 522                 } else {
 523                         current = (result < 0) ?
 524                                 current->left : current->right;
 525                 }
 526         }
 527
 528         PTHREAD_MUTEX_UNLOCK(tree);
 529         return NULL;
 530 }
 531
 532 /** Find the user data.
 533  *
 534  */
 535 void *rbtree_finddata(rbtree_t *tree, void const *data)
 536 {
 537         rbnode_t *x;
 538
 539         x = rbtree_find(tree, data);
 540         if (!x) return NULL;
 541
 542         return x->data;
 543 }
 544
 545 /** Walk the tree, Pre-order
 546  *
 547  * We call ourselves recursively for each function, but that's OK,
 548  * as the stack is only log(N) deep, which is ~12 entries deep.
 549  */
 550 static int walk_node_pre_order(rbnode_t *x, rb_walker_t compare, void *context)
 551 {
 552         int rcode;
 553         rbnode_t *left, *right;
 554
 555         left = x->left;
 556         right = x->right;
 557
 558         rcode = compare(context, x->data);
 559         if (rcode != 0) return rcode;
 560
 561         if (left != NIL) {
 562                 rcode = walk_node_pre_order(left, compare, context);
 563                 if (rcode != 0) return rcode;
 564         }
 565
 566         if (right != NIL) {
 567                 rcode = walk_node_pre_order(right, compare, context);
 568                 if (rcode != 0) return rcode;
 569         }
 570
 571         return 0;               /* we know everything returned zero */
 572 }
 573
 574 /** rbtree_in_order
 575  *
 576  */
 577 static int walk_node_in_order(rbnode_t *x, rb_walker_t compare, void *context)
 578 {
 579         int rcode;
 580         rbnode_t *right;
 581
 582         if (x->left != NIL) {
 583                 rcode = walk_node_in_order(x->left, compare, context);
 584                 if (rcode != 0) return rcode;
 585         }
 586
 587         right = x->right;
 588
 589         rcode = compare(context, x->data);
 590         if (rcode != 0) return rcode;
 591
 592         if (right != NIL) {
 593                 rcode = walk_node_in_order(right, compare, context);
 594                 if (rcode != 0) return rcode;
 595         }
 596
 597         return 0;               /* we know everything returned zero */
 598 }
 599
 600
 601 /** rbtree_post_order
 602  *
 603  */
 604 static int walk_node_post_order(rbnode_t *x, rb_walker_t compare, void *context)
 605 {
 606         int rcode;
 607
 608         if (x->left != NIL) {
 609                 rcode = walk_node_post_order(x->left, compare, context);
 610                 if (rcode != 0) return rcode;
 611         }
 612
 613         if (x->right != NIL) {
 614                 rcode = walk_node_post_order(x->right, compare, context);
 615                 if (rcode != 0) return rcode;
 616         }
 617
 618         rcode = compare(context, x->data);
 619         if (rcode != 0) return rcode;
 620
 621         return 0;               /* we know everything returned zero */
 622 }
 623
 624
 625 /** rbtree_delete_order
 626  *
 627  *      This executes an rbtree_in_order-like walk that adapts to changes in the
 628  *      tree above it, which may occur because we allow the compare to
 629  *      tell us to delete the current node.
 630  *
 631  *      The compare should return:
 632  *
 633  *              < 0  - on error
 634  *              0    - continue walking, don't delete the node
 635  *              1    - delete the node and stop walking
 636  *              2    - delete the node and continue walking
 637  */
 638 static int walk_delete_order(rbtree_t *tree, rb_walker_t compare, void *context)
 639 {
 640         rbnode_t *solid, *x;
 641         int rcode = 0;
 642
 643         /* Keep track of last node that refused deletion. */
 644         solid = NIL;
 645         while (solid == NIL) {
 646                 x = tree->root;
 647                 if (x == NIL) break;
 648         descend:
 649                 while (x->left != NIL) {
 650                         x = x->left;
 651                 }
 652         visit:
 653                 rcode = compare(context, x->data);
 654                 if (rcode < 0) {
 655                         return rcode;
 656                 }
 657                 if (rcode) {
 658                         rbtree_delete_internal(tree, x, true);
 659                         if (rcode != 2) {
 660                                 return rcode;
 661                         }
 662                 } else {
 663                         solid = x;
 664                 }
 665         }
 666         if (solid != NIL) {
 667                 x = solid;
 668                 if (x->right != NIL) {
 669                         x = x->right;
 670                         goto descend;
 671                 }
 672                 while (x->parent) {
 673                         if (x->parent->left == x) {
 674                                 x = x->parent;
 675                                 goto visit;
 676                         }
 677                         x = x->parent;
 678                 }
 679         }
 680         return rcode;
 681 }
 682
 683
 684 /*
 685  *      walk the entire tree.  The compare function CANNOT modify
 686  *      the tree.
 687  *
 688  *      The compare function should return 0 to continue walking.
 689  *      Any other value stops the walk, and is returned.
 690  */
 691 int rbtree_walk(rbtree_t *tree, rb_order_t order, rb_walker_t compare, void *context)
 692 {
 693         int rcode;
 694
 695         if (tree->root == NIL) return 0;
 696
 697         PTHREAD_MUTEX_LOCK(tree);
 698
 699         switch (order) {
 700         case RBTREE_PRE_ORDER:
 701                 rcode = walk_node_pre_order(tree->root, compare, context);
 702                 break;
 703
 704         case RBTREE_IN_ORDER:
 705                 rcode = walk_node_in_order(tree->root, compare, context);
 706                 break;
 707
 708         case RBTREE_POST_ORDER:
 709                 rcode = walk_node_post_order(tree->root, compare, context);
 710                 break;
 711
 712         case RBTREE_DELETE_ORDER:
 713                 rcode = walk_delete_order(tree, compare, context);
 714                 break;
 715
 716         default:
 717                 rcode = -1;
 718                 break;
 719         }
 720
 721         PTHREAD_MUTEX_UNLOCK(tree);
 722         return rcode;
 723 }
 724
 725 uint32_t rbtree_num_elements(rbtree_t *tree)
 726 {
 727         if (!tree) return 0;
 728
 729         return tree->num_elements;
 730 }
 731
 732 /*
 733  *      Given a Node, return the data.
 734  */
 735 void *rbtree_node2data(UNUSED rbtree_t *tree, rbnode_t *node)
 736 {
 737         if (!node) return NULL;
 738
 739         return node->data;
 740 }