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