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 #include <freeradius-devel/ident.h>
  24 RCSID("$Id$")
  25
  26 #include <freeradius-devel/autoconf.h>
  27
  28 #include <stdlib.h>
  29 #include <string.h>
  30
  31 #include <freeradius-devel/missing.h>
  32 #include <freeradius-devel/libradius.h>
  33
  34 /* red-black tree description */
  35 typedef enum { Black, Red } NodeColor;
  36
  37 struct rbnode_t {
  38     rbnode_t    *Left;          /* left child */
  39     rbnode_t    *Right;         /* right child */
  40     rbnode_t    *Parent;        /* parent */
  41     NodeColor   Color;          /* node color (black, red) */
  42     void        *Data;          /* data stored in node */
  43 };
  44
  45 #define NIL &Sentinel           /* all leafs are sentinels */
  46 static rbnode_t Sentinel = { NIL, NIL, NULL, Black, NULL};
  47
  48 struct rbtree_t {
  49 #ifndef NDEBUG
  50         uint32_t magic;
  51 #endif
  52         rbnode_t *Root;
  53         int     num_elements;
  54         int (*Compare)(const void *, const void *);
  55         int replace_flag;
  56         void (*freeNode)(void *);
  57 };
  58 #define RBTREE_MAGIC (0x5ad09c42)
  59
  60 /*
  61  *      Walks the tree to delete all nodes.
  62  *      Does NOT re-balance it!
  63  */
  64 static void FreeWalker(rbtree_t *tree, rbnode_t *X)
  65 {
  66         if (X->Left != NIL) FreeWalker(tree, X->Left);
  67         if (X->Right != NIL) FreeWalker(tree, X->Right);
  68
  69         if (tree->freeNode) tree->freeNode(X->Data);
  70         free(X);
  71 }
  72
  73 void rbtree_free(rbtree_t *tree)
  74 {
  75         if (!tree) return;
  76
  77         /*
  78          *      Walk the tree, deleting the nodes...
  79          */
  80         if (tree->Root != NIL) FreeWalker(tree, tree->Root);
  81
  82 #ifndef NDEBUG
  83         tree->magic = 0;
  84 #endif
  85         tree->Root = NULL;
  86         free(tree);
  87 }
  88
  89 /*
  90  *      Create a new red-black tree.
  91  */
  92 rbtree_t *rbtree_create(int (*Compare)(const void *, const void *),
  93                         void (*freeNode)(void *),
  94                         int replace_flag)
  95 {
  96         rbtree_t        *tree;
  97
  98         if (!Compare) return NULL;
  99
 100         tree = malloc(sizeof(*tree));
 101         if (!tree) return NULL;
 102
 103         memset(tree, 0, sizeof(*tree));
 104 #ifndef NDEBUG
 105         tree->magic = RBTREE_MAGIC;
 106 #endif
 107         tree->Root = NIL;
 108         tree->Compare = Compare;
 109         tree->replace_flag = replace_flag;
 110         tree->freeNode = freeNode;
 111
 112         return tree;
 113 }
 114
 115
 116 static void RotateLeft(rbtree_t *tree, rbnode_t *X)
 117 {
 118         /**************************
 119          *  rotate Node X to left *
 120          **************************/
 121
 122         rbnode_t *Y = X->Right;
 123
 124         /* establish X->Right link */
 125         X->Right = Y->Left;
 126         if (Y->Left != NIL) Y->Left->Parent = X;
 127
 128         /* establish Y->Parent link */
 129         if (Y != NIL) Y->Parent = X->Parent;
 130         if (X->Parent) {
 131                 if (X == X->Parent->Left)
 132                         X->Parent->Left = Y;
 133                 else
 134                         X->Parent->Right = Y;
 135         } else {
 136                 tree->Root = Y;
 137         }
 138
 139         /* link X and Y */
 140         Y->Left = X;
 141         if (X != NIL) X->Parent = Y;
 142 }
 143
 144 static void RotateRight(rbtree_t *tree, rbnode_t *X)
 145 {
 146         /****************************
 147          *  rotate Node X to right  *
 148          ****************************/
 149
 150         rbnode_t *Y = X->Left;
 151
 152         /* establish X->Left link */
 153         X->Left = Y->Right;
 154         if (Y->Right != NIL) Y->Right->Parent = X;
 155
 156         /* establish Y->Parent link */
 157         if (Y != NIL) Y->Parent = X->Parent;
 158         if (X->Parent) {
 159                 if (X == X->Parent->Right)
 160                         X->Parent->Right = Y;
 161                 else
 162                         X->Parent->Left = Y;
 163         } else {
 164                 tree->Root = Y;
 165         }
 166
 167         /* link X and Y */
 168         Y->Right = X;
 169         if (X != NIL) X->Parent = Y;
 170 }
 171
 172 static void InsertFixup(rbtree_t *tree, rbnode_t *X)
 173 {
 174         /*************************************
 175          *  maintain red-black tree balance  *
 176          *  after inserting node X           *
 177          *************************************/
 178
 179         /* check red-black properties */
 180         while (X != tree->Root && X->Parent->Color == Red) {
 181                 /* we have a violation */
 182                 if (X->Parent == X->Parent->Parent->Left) {
 183                         rbnode_t *Y = X->Parent->Parent->Right;
 184                         if (Y->Color == Red) {
 185
 186                                 /* uncle is red */
 187                                 X->Parent->Color = Black;
 188                                 Y->Color = Black;
 189                                 X->Parent->Parent->Color = Red;
 190                                 X = X->Parent->Parent;
 191                         } else {
 192
 193                                 /* uncle is black */
 194                                 if (X == X->Parent->Right) {
 195                                         /* make X a left child */
 196                                         X = X->Parent;
 197                                         RotateLeft(tree, X);
 198                                 }
 199
 200                                 /* recolor and rotate */
 201                                 X->Parent->Color = Black;
 202                                 X->Parent->Parent->Color = Red;
 203                                 RotateRight(tree, X->Parent->Parent);
 204                         }
 205                 } else {
 206
 207                         /* mirror image of above code */
 208                         rbnode_t *Y = X->Parent->Parent->Left;
 209                         if (Y->Color == Red) {
 210
 211                                 /* uncle is red */
 212                                 X->Parent->Color = Black;
 213                                 Y->Color = Black;
 214                                 X->Parent->Parent->Color = Red;
 215                                 X = X->Parent->Parent;
 216                         } else {
 217
 218                                 /* uncle is black */
 219                                 if (X == X->Parent->Left) {
 220                                         X = X->Parent;
 221                                         RotateRight(tree, X);
 222                                 }
 223                                 X->Parent->Color = Black;
 224                                 X->Parent->Parent->Color = Red;
 225                                 RotateLeft(tree, X->Parent->Parent);
 226                         }
 227                 }
 228         }
 229
 230         tree->Root->Color = Black;
 231 }
 232
 233
 234 /*
 235  *      Insert an element into the tree.
 236  */
 237 int rbtree_insert(rbtree_t *tree, void *Data)
 238 {
 239         rbnode_t *Current, *Parent, *X;
 240
 241         /***********************************************
 242          *  allocate node for Data and insert in tree  *
 243          ***********************************************/
 244
 245         /* find where node belongs */
 246         Current = tree->Root;
 247         Parent = NULL;
 248         while (Current != NIL) {
 249                 int result;
 250
 251                 /*
 252                  *      See if two entries are identical.
 253                  */
 254                 result = tree->Compare(Data, Current->Data);
 255                 if (result == 0) {
 256                         /*
 257                          *      Don't replace the entry.
 258                          */
 259                         if (tree->replace_flag == 0) {
 260                                 return 0;
 261                         }
 262
 263                         /*
 264                          *      Do replace the entry.
 265                          */
 266                         if (tree->freeNode) tree->freeNode(Current->Data);
 267                         Current->Data = Data;
 268                         return 1;
 269                 }
 270
 271                 Parent = Current;
 272                 Current = (result < 0) ? Current->Left : Current->Right;
 273         }
 274
 275         /* setup new node */
 276         if ((X = malloc (sizeof(*X))) == NULL) {
 277                 exit(1);        /* FIXME! */
 278         }
 279
 280         X->Data = Data;
 281         X->Parent = Parent;
 282         X->Left = NIL;
 283         X->Right = NIL;
 284         X->Color = Red;
 285
 286         /* insert node in tree */
 287         if (Parent) {
 288                 if (tree->Compare(Data, Parent->Data) <= 0)
 289                         Parent->Left = X;
 290                 else
 291                         Parent->Right = X;
 292         } else {
 293                 tree->Root = X;
 294         }
 295
 296         InsertFixup(tree, X);
 297
 298         tree->num_elements++;
 299
 300         return 1;
 301 }
 302
 303 static void DeleteFixup(rbtree_t *tree, rbnode_t *X, rbnode_t *Parent)
 304 {
 305         /*************************************
 306          *  maintain red-black tree balance  *
 307          *  after deleting node X            *
 308          *************************************/
 309
 310         while (X != tree->Root && X->Color == Black) {
 311                 if (X == Parent->Left) {
 312                         rbnode_t *W = Parent->Right;
 313                         if (W->Color == Red) {
 314                                 W->Color = Black;
 315                                 Parent->Color = Red; /* Parent != NIL? */
 316                                 RotateLeft(tree, Parent);
 317                                 W = Parent->Right;
 318                         }
 319                         if (W->Left->Color == Black && W->Right->Color == Black) {
 320                                 if (W != NIL) W->Color = Red;
 321                                 X = Parent;
 322                                 Parent = X->Parent;
 323                         } else {
 324                                 if (W->Right->Color == Black) {
 325                                         if (W->Left != NIL) W->Left->Color = Black;
 326                                         W->Color = Red;
 327                                         RotateRight(tree, W);
 328                                         W = Parent->Right;
 329                                 }
 330                                 W->Color = Parent->Color;
 331                                 if (Parent != NIL) Parent->Color = Black;
 332                                 if (W->Right->Color != Black) {
 333                                         W->Right->Color = Black;
 334                                 }
 335                                 RotateLeft(tree, Parent);
 336                                 X = tree->Root;
 337                         }
 338                 } else {
 339                         rbnode_t *W = Parent->Left;
 340                         if (W->Color == Red) {
 341                                 W->Color = Black;
 342                                 Parent->Color = Red; /* Parent != NIL? */
 343                                 RotateRight(tree, Parent);
 344                                 W = Parent->Left;
 345                         }
 346                         if (W->Right->Color == Black && W->Left->Color == Black) {
 347                                 if (W != NIL) W->Color = Red;
 348                                 X = Parent;
 349                                 Parent = X->Parent;
 350                         } else {
 351                                 if (W->Left->Color == Black) {
 352                                         if (W->Right != NIL) W->Right->Color = Black;
 353                                         W->Color = Red;
 354                                         RotateLeft(tree, W);
 355                                         W = Parent->Left;
 356                                 }
 357                                 W->Color = Parent->Color;
 358                                 if (Parent != NIL) Parent->Color = Black;
 359                                 if (W->Left->Color != Black) {
 360                                         W->Left->Color = Black;
 361                                 }
 362                                 RotateRight(tree, Parent);
 363                                 X = tree->Root;
 364                         }
 365                 }
 366         }
 367         X->Color = Black;
 368 }
 369
 370 /*
 371  *      Delete an element from the tree.
 372  */
 373 void rbtree_delete(rbtree_t *tree, rbnode_t *Z)
 374 {
 375         rbnode_t *X, *Y;
 376         rbnode_t *Parent;
 377
 378         /*****************************
 379          *  delete node Z from tree  *
 380          *****************************/
 381
 382         if (!Z || Z == NIL) return;
 383
 384         if (Z->Left == NIL || Z->Right == NIL) {
 385                 /* Y has a NIL node as a child */
 386                 Y = Z;
 387         } else {
 388                 /* find tree successor with a NIL node as a child */
 389                 Y = Z->Right;
 390                 while (Y->Left != NIL) Y = Y->Left;
 391         }
 392
 393         /* X is Y's only child */
 394         if (Y->Left != NIL)
 395                 X = Y->Left;
 396         else
 397                 X = Y->Right;   /* may be NIL! */
 398
 399         /* remove Y from the parent chain */
 400         Parent = Y->Parent;
 401         if (X != NIL) X->Parent = Parent;
 402
 403         if (Parent)
 404                 if (Y == Parent->Left)
 405                         Parent->Left = X;
 406                 else
 407                         Parent->Right = X;
 408         else
 409                 tree->Root = X;
 410
 411         if (Y != Z) {
 412                 /*
 413                  *      Move the child's data to here, and then
 414                  *      re-balance the tree.
 415                  */
 416                 if (tree->freeNode) tree->freeNode(Z->Data);
 417                 Z->Data = Y->Data;
 418                 Y->Data = NULL;
 419         } else if (tree->freeNode) {
 420                 tree->freeNode(Z->Data);
 421         }
 422
 423         if (Y->Color == Black && X != NIL)
 424                 DeleteFixup(tree, X, Parent);
 425
 426         free(Y);
 427
 428         tree->num_elements--;
 429 }
 430
 431 /*
 432  *      Delete a node from the tree, based on given data, which MUST
 433  *      have come from rbtree_finddata().
 434  */
 435 int rbtree_deletebydata(rbtree_t *tree, const void *data)
 436 {
 437         rbnode_t *node = rbtree_find(tree, data);
 438
 439         if (!node) return 0;    /* false */
 440
 441         rbtree_delete(tree, node);
 442
 443         return 1;
 444 }
 445
 446
 447 /*
 448  *      Find an element in the tree, returning the data, not the node.
 449  */
 450 rbnode_t *rbtree_find(rbtree_t *tree, const void *Data)
 451 {
 452         /*******************************
 453          *  find node containing Data  *
 454          *******************************/
 455
 456         rbnode_t *Current = tree->Root;
 457
 458         while (Current != NIL) {
 459                 int result = tree->Compare(Data, Current->Data);
 460
 461                 if (result == 0) {
 462                         return Current;
 463                 } else {
 464                         Current = (result < 0) ?
 465                                 Current->Left : Current->Right;
 466                 }
 467         }
 468         return NULL;
 469 }
 470
 471 /*
 472  *      Find the user data.
 473  */
 474 void *rbtree_finddata(rbtree_t *tree, const void *Data)
 475 {
 476         rbnode_t *X;
 477
 478         X = rbtree_find(tree, Data);
 479         if (!X) return NULL;
 480
 481         return X->Data;
 482 }
 483
 484 /*
 485  *      Walk the tree, Pre-order
 486  *
 487  *      We call ourselves recursively for each function, but that's OK,
 488  *      as the stack is only log(N) deep, which is ~12 entries deep.
 489  */
 490 static int WalkNodePreOrder(rbnode_t *X,
 491                             int (*callback)(void *, void *), void *context)
 492 {
 493         int rcode;
 494
 495         rcode = callback(context, X->Data);
 496         if (rcode != 0) return rcode;
 497
 498         if (X->Left != NIL) {
 499                 rcode = WalkNodePreOrder(X->Left, callback, context);
 500                 if (rcode != 0) return rcode;
 501         }
 502
 503         if (X->Right != NIL) {
 504                 rcode = WalkNodePreOrder(X->Right, callback, context);
 505                 if (rcode != 0) return rcode;
 506         }
 507
 508         return 0;               /* we know everything returned zero */
 509 }
 510
 511 /*
 512  *      Inorder
 513  */
 514 static int WalkNodeInOrder(rbnode_t *X,
 515                            int (*callback)(void *, void *), void *context)
 516 {
 517         int rcode;
 518
 519         if (X->Left != NIL) {
 520                 rcode = WalkNodeInOrder(X->Left, callback, context);
 521                 if (rcode != 0) return rcode;
 522         }
 523
 524         rcode = callback(context, X->Data);
 525         if (rcode != 0) return rcode;
 526
 527         if (X->Right != NIL) {
 528                 rcode = WalkNodeInOrder(X->Right, callback, context);
 529                 if (rcode != 0) return rcode;
 530         }
 531
 532         return 0;               /* we know everything returned zero */
 533 }
 534
 535
 536 /*
 537  *      PostOrder
 538  */
 539 static int WalkNodePostOrder(rbnode_t *X,
 540                              int (*callback)(void *, void*), void *context)
 541 {
 542         int rcode;
 543
 544         if (X->Left != NIL) {
 545                 rcode = WalkNodeInOrder(X->Left, callback, context);
 546                 if (rcode != 0) return rcode;
 547         }
 548
 549         if (X->Right != NIL) {
 550                 rcode = WalkNodeInOrder(X->Right, callback, context);
 551                 if (rcode != 0) return rcode;
 552         }
 553
 554         rcode = callback(context, X->Data);
 555         if (rcode != 0) return rcode;
 556
 557         return 0;               /* we know everything returned zero */
 558 }
 559
 560 /*
 561  *      Walk the entire tree.  The callback function CANNOT modify
 562  *      the tree.
 563  *
 564  *      The callback function should return 0 to continue walking.
 565  *      Any other value stops the walk, and is returned.
 566  */
 567 int rbtree_walk(rbtree_t *tree, RBTREE_ORDER order,
 568                 int (*callback)(void *, void *), void *context)
 569 {
 570         switch (order) {
 571         case PreOrder:
 572                 return WalkNodePreOrder(tree->Root, callback, context);
 573         case InOrder:
 574                 return WalkNodeInOrder(tree->Root, callback, context);
 575         case PostOrder:
 576                 return WalkNodePostOrder(tree->Root, callback, context);
 577
 578         default:
 579                 break;
 580         }
 581
 582         return -1;
 583 }
 584
 585 int rbtree_num_elements(rbtree_t *tree)
 586 {
 587         if (!tree) return 0;
 588
 589         return tree->num_elements;
 590 }
 591
 592
 593 /*
 594  *      Given a Node, return the data.
 595  */
 596 void *rbtree_node2data(rbtree_t *tree, rbnode_t *node)
 597 {
 598         tree = tree;            /* -Wunused */
 599
 600         if (!node) return NULL;
 601
 602         return node->Data;
 603 }