src/lib/rbtree.c

   1 /*
   2  * rbtree.c     Red-black balanced binary trees.
   3  *
   4  * Version:     $Id$
   5  *
   6  *   This library is free software; you can redistribute it and/or
   7  *   modify it under the terms of the GNU Lesser General Public
   8  *   License as published by the Free Software Foundation; either
   9  *   version 2.1 of the License, or (at your option) any later version.
  10  *
  11  *   This library 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 GNU
  14  *   Lesser General Public License for more details.
  15  *
  16  *   You should have received a copy of the GNU Lesser General Public
  17  *   License along with this library; 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 { Black, Red } NodeColor;
  39
  40 struct rbnode_t {
  41     rbnode_t    *Left;          /* left child */
  42     rbnode_t    *Right;         /* right child */
  43     rbnode_t    *Parent;        /* parent */
  44     NodeColor   Color;          /* node color (black, red) */
  45     void        *Data;          /* data stored in node */
  46 };
  47
  48 #define NIL &Sentinel      /* all leafs are sentinels */
  49 static rbnode_t Sentinel = { NIL, NIL, NULL, Black, NULL};
  50
  51 struct rbtree_t {
  52 #ifndef NDEBUG
  53         uint32_t magic;
  54 #endif
  55         rbnode_t *Root;
  56         int     num_elements;
  57         int (*Compare)(void const *, void const *);
  58         int replace_flag;
  59         void (*freeNode)(void *);
  60 #ifdef HAVE_PTHREAD_H
  61         int lock;
  62         pthread_mutex_t mutex;
  63 #endif
  64 };
  65 #define RBTREE_MAGIC (0x5ad09c42)
  66
  67 /*
  68  *      Walks the tree to delete all nodes.
  69  *      Does NOT re-balance it!
  70  */
  71 static void FreeWalker(rbtree_t *tree, rbnode_t *X)
  72 {
  73         if (X->Left != NIL) FreeWalker(tree, X->Left);
  74         if (X->Right != NIL) FreeWalker(tree, X->Right);
  75
  76         if (tree->freeNode) tree->freeNode(X->Data);
  77         free(X);
  78 }
  79
  80 void rbtree_free(rbtree_t *tree)
  81 {
  82         if (!tree) return;
  83
  84         PTHREAD_MUTEX_LOCK(tree);
  85
  86         /*
  87          *      Walk the tree, deleting the nodes...
  88          */
  89         if (tree->Root != NIL) FreeWalker(tree, tree->Root);
  90
  91 #ifndef NDEBUG
  92         tree->magic = 0;
  93 #endif
  94         tree->Root = NULL;
  95
  96 #ifdef HAVE_PTHREAD_H
  97         if (tree->lock) pthread_mutex_destroy(&tree->mutex);
  98 #endif
  99
 100         free(tree);
 101 }
 102
 103 /*
 104  *      Create a new red-black tree.
 105  */
 106 rbtree_t *rbtree_create(int (*Compare)(void const *, void const *),
 107                         void (*freeNode)(void *),
 108                         int flags)
 109 {
 110         rbtree_t        *tree;
 111
 112         if (!Compare) return NULL;
 113
 114         tree = malloc(sizeof(*tree));
 115         if (!tree) return NULL;
 116
 117         memset(tree, 0, sizeof(*tree));
 118 #ifndef NDEBUG
 119         tree->magic = RBTREE_MAGIC;
 120 #endif
 121         tree->Root = NIL;
 122         tree->Compare = Compare;
 123         tree->replace_flag = flags & RBTREE_FLAG_REPLACE;
 124 #ifdef HAVE_PTHREAD_H
 125         tree->lock = flags & RBTREE_FLAG_LOCK;
 126         if (tree->lock) {
 127                 pthread_mutex_init(&tree->mutex, NULL);
 128         }
 129 #endif
 130         tree->freeNode = freeNode;
 131
 132         return tree;
 133 }
 134
 135
 136 static void RotateLeft(rbtree_t *tree, rbnode_t *X)
 137 {
 138         /**************************
 139          *  rotate Node X to left *
 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         } else {
 156                 tree->Root = Y;
 157         }
 158
 159         /* link X and Y */
 160         Y->Left = X;
 161         if (X != NIL) X->Parent = Y;
 162 }
 163
 164 static void RotateRight(rbtree_t *tree, rbnode_t *X)
 165 {
 166         /****************************
 167          *  rotate Node X to right  *
 168          ****************************/
 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         } else {
 184                 tree->Root = Y;
 185         }
 186
 187         /* link X and Y */
 188         Y->Right = X;
 189         if (X != NIL) X->Parent = Y;
 190 }
 191
 192 static void InsertFixup(rbtree_t *tree, rbnode_t *X)
 193 {
 194         /*************************************
 195          *  maintain red-black tree balance  *
 196          *  after inserting node X         *
 197          *************************************/
 198
 199         /* check red-black properties */
 200         while (X != tree->Root && X->Parent->Color == Red) {
 201                 /* we have a violation */
 202                 if (X->Parent == X->Parent->Parent->Left) {
 203                         rbnode_t *Y = X->Parent->Parent->Right;
 204                         if (Y->Color == Red) {
 205
 206                                 /* uncle is red */
 207                                 X->Parent->Color = Black;
 208                                 Y->Color = Black;
 209                                 X->Parent->Parent->Color = Red;
 210                                 X = X->Parent->Parent;
 211                         } else {
 212
 213                                 /* uncle is black */
 214                                 if (X == X->Parent->Right) {
 215                                         /* make X a left child */
 216                                         X = X->Parent;
 217                                         RotateLeft(tree, X);
 218                                 }
 219
 220                                 /* recolor and rotate */
 221                                 X->Parent->Color = Black;
 222                                 X->Parent->Parent->Color = Red;
 223                                 RotateRight(tree, X->Parent->Parent);
 224                         }
 225                 } else {
 226
 227                         /* mirror image of above code */
 228                         rbnode_t *Y = X->Parent->Parent->Left;
 229                         if (Y->Color == Red) {
 230
 231                                 /* uncle is red */
 232                                 X->Parent->Color = Black;
 233                                 Y->Color = Black;
 234                                 X->Parent->Parent->Color = Red;
 235                                 X = X->Parent->Parent;
 236                         } else {
 237
 238                                 /* uncle is black */
 239                                 if (X == X->Parent->Left) {
 240                                         X = X->Parent;
 241                                         RotateRight(tree, X);
 242                                 }
 243                                 X->Parent->Color = Black;
 244                                 X->Parent->Parent->Color = Red;
 245                                 RotateLeft(tree, X->Parent->Parent);
 246                         }
 247                 }
 248         }
 249
 250         tree->Root->Color = Black;
 251 }
 252
 253
 254 /*
 255  *      Insert an element into the tree.
 256  */
 257 rbnode_t *rbtree_insertnode(rbtree_t *tree, void *Data)
 258 {
 259         rbnode_t *Current, *Parent, *X;
 260
 261         /***********************************************
 262          *  allocate node for Data and insert in tree  *
 263          ***********************************************/
 264
 265         PTHREAD_MUTEX_LOCK(tree);
 266
 267         /* find where node belongs */
 268         Current = tree->Root;
 269         Parent = NULL;
 270         while (Current != NIL) {
 271                 int result;
 272
 273                 /*
 274                  *      See if two entries are identical.
 275                  */
 276                 result = tree->Compare(Data, Current->Data);
 277                 if (result == 0) {
 278                         /*
 279                          *      Don't replace the entry.
 280                          */
 281                         if (tree->replace_flag == 0) {
 282                                 PTHREAD_MUTEX_UNLOCK(tree);
 283                                 return NULL;
 284                         }
 285
 286                         /*
 287                          *      Do replace the entry.
 288                          */
 289                         if (tree->freeNode) tree->freeNode(Current->Data);
 290                         Current->Data = Data;
 291                         PTHREAD_MUTEX_UNLOCK(tree);
 292                         return Current;
 293                 }
 294
 295                 Parent = Current;
 296                 Current = (result < 0) ? Current->Left : Current->Right;
 297         }
 298
 299         /* setup new node */
 300         if ((X = malloc (sizeof(*X))) == NULL) {
 301                 exit(1);        /* FIXME! */
 302         }
 303
 304         X->Data = Data;
 305         X->Parent = Parent;
 306         X->Left = NIL;
 307         X->Right = NIL;
 308         X->Color = Red;
 309
 310         /* insert node in tree */
 311         if (Parent) {
 312                 if (tree->Compare(Data, Parent->Data) <= 0)
 313                         Parent->Left = X;
 314                 else
 315                         Parent->Right = X;
 316         } else {
 317                 tree->Root = X;
 318         }
 319
 320         InsertFixup(tree, X);
 321
 322         tree->num_elements++;
 323
 324         PTHREAD_MUTEX_UNLOCK(tree);
 325         return X;
 326 }
 327
 328 bool rbtree_insert(rbtree_t *tree, void *Data)
 329 {
 330         if (rbtree_insertnode(tree, Data)) return true;
 331         return false;
 332 }
 333
 334 static void DeleteFixup(rbtree_t *tree, rbnode_t *X, rbnode_t *Parent)
 335 {
 336         /*************************************
 337          *  maintain red-black tree balance  *
 338          *  after deleting node X           *
 339          *************************************/
 340
 341         while (X != tree->Root && X->Color == Black) {
 342                 if (X == Parent->Left) {
 343                         rbnode_t *W = Parent->Right;
 344                         if (W->Color == Red) {
 345                                 W->Color = Black;
 346                                 Parent->Color = Red; /* Parent != NIL? */
 347                                 RotateLeft(tree, Parent);
 348                                 W = Parent->Right;
 349                         }
 350                         if (W->Left->Color == Black && W->Right->Color == Black) {
 351                                 if (W != NIL) W->Color = Red;
 352                                 X = Parent;
 353                                 Parent = X->Parent;
 354                         } else {
 355                                 if (W->Right->Color == Black) {
 356                                         if (W->Left != NIL) W->Left->Color = Black;
 357                                         W->Color = Red;
 358                                         RotateRight(tree, W);
 359                                         W = Parent->Right;
 360                                 }
 361                                 W->Color = Parent->Color;
 362                                 if (Parent != NIL) Parent->Color = Black;
 363                                 if (W->Right->Color != Black) {
 364                                         W->Right->Color = Black;
 365                                 }
 366                                 RotateLeft(tree, Parent);
 367                                 X = tree->Root;
 368                         }
 369                 } else {
 370                         rbnode_t *W = Parent->Left;
 371                         if (W->Color == Red) {
 372                                 W->Color = Black;
 373                                 Parent->Color = Red; /* Parent != NIL? */
 374                                 RotateRight(tree, Parent);
 375                                 W = Parent->Left;
 376                         }
 377                         if (W->Right->Color == Black && W->Left->Color == Black) {
 378                                 if (W != NIL) W->Color = Red;
 379                                 X = Parent;
 380                                 Parent = X->Parent;
 381                         } else {
 382                                 if (W->Left->Color == Black) {
 383                                         if (W->Right != NIL) W->Right->Color = Black;
 384                                         W->Color = Red;
 385                                         RotateLeft(tree, W);
 386                                         W = Parent->Left;
 387                                 }
 388                                 W->Color = Parent->Color;
 389                                 if (Parent != NIL) Parent->Color = Black;
 390                                 if (W->Left->Color != Black) {
 391                                         W->Left->Color = Black;
 392                                 }
 393                                 RotateRight(tree, Parent);
 394                                 X = tree->Root;
 395                         }
 396                 }
 397         }
 398         if (X != NIL) X->Color = Black; /* Avoid cache-dirty on NIL */
 399 }
 400
 401 /*
 402  *      Delete an element from the tree.
 403  */
 404 void rbtree_delete(rbtree_t *tree, rbnode_t *Z)
 405 {
 406         rbnode_t *X, *Y;
 407         rbnode_t *Parent;
 408
 409         /*****************************
 410          *  delete node Z from tree  *
 411          *****************************/
 412
 413         if (!Z || Z == NIL) return;
 414
 415         PTHREAD_MUTEX_LOCK(tree);
 416
 417         if (Z->Left == NIL || Z->Right == NIL) {
 418                 /* Y has a NIL node as a child */
 419                 Y = Z;
 420         } else {
 421                 /* find tree successor with a NIL node as a child */
 422                 Y = Z->Right;
 423                 while (Y->Left != NIL) Y = Y->Left;
 424         }
 425
 426         /* X is Y's only child */
 427         if (Y->Left != NIL)
 428                 X = Y->Left;
 429         else
 430                 X = Y->Right;   /* may be NIL! */
 431
 432         /* remove Y from the parent chain */
 433         Parent = Y->Parent;
 434         if (X != NIL) X->Parent = Parent;
 435
 436         if (Parent)
 437                 if (Y == Parent->Left)
 438                         Parent->Left = X;
 439                 else
 440                         Parent->Right = X;
 441         else
 442                 tree->Root = X;
 443
 444         if (Y != Z) {
 445                 if (tree->freeNode) tree->freeNode(Z->Data);
 446                 Z->Data = Y->Data;
 447                 Y->Data = NULL;
 448
 449                 if (Y->Color == Black)
 450                         DeleteFixup(tree, X, Parent);
 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                 free(Z);
 471
 472         } else {
 473                 if (tree->freeNode) tree->freeNode(Y->Data);
 474
 475                 if (Y->Color == Black)
 476                         DeleteFixup(tree, X, Parent);
 477
 478                 free(Y);
 479         }
 480
 481         tree->num_elements--;
 482         PTHREAD_MUTEX_UNLOCK(tree);
 483 }
 484
 485 /*
 486  *      Delete a node from the tree, based on given data, which MUST
 487  *      have come from rbtree_finddata().
 488  */
 489 bool rbtree_deletebydata(rbtree_t *tree, void const *data)
 490 {
 491         rbnode_t *node = rbtree_find(tree, data);
 492
 493         if (!node) return false;
 494
 495         rbtree_delete(tree, node);
 496
 497         return true;
 498 }
 499
 500
 501 /*
 502  *      Find an element in the tree, returning the data, not the node.
 503  */
 504 rbnode_t *rbtree_find(rbtree_t *tree, void const *Data)
 505 {
 506         /*******************************
 507          *  find node containing Data  *
 508          *******************************/
 509
 510         rbnode_t *Current;
 511
 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 /*
 533  *      Find the user data.
 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 /*
 546  *      Walk the tree, Pre-order
 547  *
 548  *      We call ourselves recursively for each function, but that's OK,
 549  *      as the stack is only log(N) deep, which is ~12 entries deep.
 550  */
 551 static int WalkNodePreOrder(rbnode_t *X,
 552                             int (*callback)(void *, void *), void *context)
 553 {
 554         int rcode;
 555         rbnode_t *Left, *Right;
 556
 557         Left = X->Left;
 558         Right = X->Right;
 559
 560         rcode = callback(context, X->Data);
 561         if (rcode != 0) return rcode;
 562
 563         if (Left != NIL) {
 564                 rcode = WalkNodePreOrder(Left, callback, context);
 565                 if (rcode != 0) return rcode;
 566         }
 567
 568         if (Right != NIL) {
 569                 rcode = WalkNodePreOrder(Right, callback, context);
 570                 if (rcode != 0) return rcode;
 571         }
 572
 573         return 0;               /* we know everything returned zero */
 574 }
 575
 576 /*
 577  *      Inorder
 578  */
 579 static int WalkNodeInOrder(rbnode_t *X,
 580                            int (*callback)(void *, void *), void *context)
 581 {
 582         int rcode;
 583         rbnode_t *Right;
 584
 585         if (X->Left != NIL) {
 586                 rcode = WalkNodeInOrder(X->Left, callback, context);
 587                 if (rcode != 0) return rcode;
 588         }
 589
 590         Right = X->Right;
 591
 592         rcode = callback(context, X->Data);
 593         if (rcode != 0) return rcode;
 594
 595         if (Right != NIL) {
 596                 rcode = WalkNodeInOrder(Right, callback, context);
 597                 if (rcode != 0) return rcode;
 598         }
 599
 600         return 0;               /* we know everything returned zero */
 601 }
 602
 603
 604 /*
 605  *      PostOrder
 606  */
 607 static int WalkNodePostOrder(rbnode_t *X,
 608                              int (*callback)(void *, void*), void *context)
 609 {
 610         int rcode;
 611
 612         if (X->Left != NIL) {
 613                 rcode = WalkNodePostOrder(X->Left, callback, context);
 614                 if (rcode != 0) return rcode;
 615         }
 616
 617         if (X->Right != NIL) {
 618                 rcode = WalkNodePostOrder(X->Right, callback, context);
 619                 if (rcode != 0) return rcode;
 620         }
 621
 622         rcode = callback(context, X->Data);
 623         if (rcode != 0) return rcode;
 624
 625         return 0;               /* we know everything returned zero */
 626 }
 627
 628 /*
 629  *      Walk the entire tree.  The callback function CANNOT modify
 630  *      the tree.
 631  *
 632  *      The callback function should return 0 to continue walking.
 633  *      Any other value stops the walk, and is returned.
 634  */
 635 int rbtree_walk(rbtree_t *tree, RBTREE_ORDER order,
 636                 int (*callback)(void *, void *), void *context)
 637 {
 638         int rcode;
 639
 640         if (tree->Root == NIL) return 0;
 641
 642         PTHREAD_MUTEX_LOCK(tree);
 643         switch (order) {
 644         case PreOrder:
 645                 rcode = WalkNodePreOrder(tree->Root, callback, context);
 646                 break;
 647         case InOrder:
 648                 rcode = WalkNodeInOrder(tree->Root, callback, context);
 649                 break;
 650         case PostOrder:
 651                 rcode = WalkNodePostOrder(tree->Root, callback, context);
 652                 break;
 653         default:
 654                 rcode = -1;
 655                 break;
 656         }
 657
 658         PTHREAD_MUTEX_UNLOCK(tree);
 659         return rcode;
 660 }
 661
 662 int rbtree_num_elements(rbtree_t *tree)
 663 {
 664         if (!tree) return 0;
 665
 666         return tree->num_elements;
 667 }
 668
 669
 670 /*
 671  *      Given a Node, return the data.
 672  */
 673 void *rbtree_node2data(UNUSED rbtree_t *tree, rbnode_t *node)
 674 {
 675         if (!node) return NULL;
 676
 677         return node->Data;
 678 }
 679
 680 /*
 681  *      Return left-most child.
 682  */
 683 void *rbtree_min(rbtree_t *tree)
 684 {
 685         void *data;
 686         rbnode_t *Current;
 687
 688         if (!tree || !tree->Root) return NULL;
 689
 690         PTHREAD_MUTEX_LOCK(tree);
 691         Current = tree->Root;
 692         while (Current->Left != NIL) Current = Current->Left;
 693
 694         data = Current->Data;
 695         PTHREAD_MUTEX_UNLOCK(tree);
 696         return data;
 697 }