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