lib/typerb.c

   1 // SPDX-License-Identifier: ISC AND BSD-2-Clause
   2 /* RB-tree */
   3
   4 /*
   5  * Copyright 2002 Niels Provos <provos@citi.umich.edu>
   6  */
   7
   8 /*
   9  * Copyright (c) 2016 David Gwynne <dlg@openbsd.org>
  10  */
  11
  12 #ifdef HAVE_CONFIG_H
  13 #include "config.h"
  14 #endif
  15
  16 #include <string.h>
  17 #include "typerb.h"
  18
  19 #define RB_BLACK        0
  20 #define RB_RED          1
  21
  22 #define rb_entry                typed_rb_entry
  23 #define rbt_tree                typed_rb_root
  24
  25 #define RBE_LEFT(_rbe)          (_rbe)->rbt_left
  26 #define RBE_RIGHT(_rbe)         (_rbe)->rbt_right
  27 #define RBE_PARENT(_rbe)        (_rbe)->rbt_parent
  28 #define RBE_COLOR(_rbe)         (_rbe)->rbt_color
  29
  30 #define RBH_ROOT(_rbt)          (_rbt)->rbt_root
  31
  32 static inline void rbe_set(struct rb_entry *rbe, struct rb_entry *parent)
  33 {
  34         RBE_PARENT(rbe) = parent;
  35         RBE_LEFT(rbe) = RBE_RIGHT(rbe) = NULL;
  36         RBE_COLOR(rbe) = RB_RED;
  37 }
  38
  39 static inline void rbe_set_blackred(struct rb_entry *black,
  40                                     struct rb_entry *red)
  41 {
  42         RBE_COLOR(black) = RB_BLACK;
  43         RBE_COLOR(red) = RB_RED;
  44 }
  45
  46 static inline void rbe_rotate_left(struct rbt_tree *rbt, struct rb_entry *rbe)
  47 {
  48         struct rb_entry *parent;
  49         struct rb_entry *tmp;
  50
  51         tmp = RBE_RIGHT(rbe);
  52         RBE_RIGHT(rbe) = RBE_LEFT(tmp);
  53         if (RBE_RIGHT(rbe) != NULL)
  54                 RBE_PARENT(RBE_LEFT(tmp)) = rbe;
  55
  56         parent = RBE_PARENT(rbe);
  57         RBE_PARENT(tmp) = parent;
  58         if (parent != NULL) {
  59                 if (rbe == RBE_LEFT(parent))
  60                         RBE_LEFT(parent) = tmp;
  61                 else
  62                         RBE_RIGHT(parent) = tmp;
  63         } else
  64                 RBH_ROOT(rbt) = tmp;
  65
  66         RBE_LEFT(tmp) = rbe;
  67         RBE_PARENT(rbe) = tmp;
  68 }
  69
  70 static inline void rbe_rotate_right(struct rbt_tree *rbt, struct rb_entry *rbe)
  71 {
  72         struct rb_entry *parent;
  73         struct rb_entry *tmp;
  74
  75         tmp = RBE_LEFT(rbe);
  76         RBE_LEFT(rbe) = RBE_RIGHT(tmp);
  77         if (RBE_LEFT(rbe) != NULL)
  78                 RBE_PARENT(RBE_RIGHT(tmp)) = rbe;
  79
  80         parent = RBE_PARENT(rbe);
  81         RBE_PARENT(tmp) = parent;
  82         if (parent != NULL) {
  83                 if (rbe == RBE_LEFT(parent))
  84                         RBE_LEFT(parent) = tmp;
  85                 else
  86                         RBE_RIGHT(parent) = tmp;
  87         } else
  88                 RBH_ROOT(rbt) = tmp;
  89
  90         RBE_RIGHT(tmp) = rbe;
  91         RBE_PARENT(rbe) = tmp;
  92 }
  93
  94 static inline void rbe_insert_color(struct rbt_tree *rbt, struct rb_entry *rbe)
  95 {
  96         struct rb_entry *parent, *gparent, *tmp;
  97
  98         rbt->count++;
  99
 100         while ((parent = RBE_PARENT(rbe)) != NULL
 101                && RBE_COLOR(parent) == RB_RED) {
 102                 gparent = RBE_PARENT(parent);
 103
 104                 if (parent == RBE_LEFT(gparent)) {
 105                         tmp = RBE_RIGHT(gparent);
 106                         if (tmp != NULL && RBE_COLOR(tmp) == RB_RED) {
 107                                 RBE_COLOR(tmp) = RB_BLACK;
 108                                 rbe_set_blackred(parent, gparent);
 109                                 rbe = gparent;
 110                                 continue;
 111                         }
 112
 113                         if (RBE_RIGHT(parent) == rbe) {
 114                                 rbe_rotate_left(rbt, parent);
 115                                 tmp = parent;
 116                                 parent = rbe;
 117                                 rbe = tmp;
 118                         }
 119
 120                         rbe_set_blackred(parent, gparent);
 121                         rbe_rotate_right(rbt, gparent);
 122                 } else {
 123                         tmp = RBE_LEFT(gparent);
 124                         if (tmp != NULL && RBE_COLOR(tmp) == RB_RED) {
 125                                 RBE_COLOR(tmp) = RB_BLACK;
 126                                 rbe_set_blackred(parent, gparent);
 127                                 rbe = gparent;
 128                                 continue;
 129                         }
 130
 131                         if (RBE_LEFT(parent) == rbe) {
 132                                 rbe_rotate_right(rbt, parent);
 133                                 tmp = parent;
 134                                 parent = rbe;
 135                                 rbe = tmp;
 136                         }
 137
 138                         rbe_set_blackred(parent, gparent);
 139                         rbe_rotate_left(rbt, gparent);
 140                 }
 141         }
 142
 143         RBE_COLOR(RBH_ROOT(rbt)) = RB_BLACK;
 144 }
 145
 146 static inline void rbe_remove_color(struct rbt_tree *rbt,
 147                                     struct rb_entry *parent,
 148                                     struct rb_entry *rbe)
 149 {
 150         struct rb_entry *tmp;
 151
 152         while ((rbe == NULL || RBE_COLOR(rbe) == RB_BLACK)
 153                && rbe != RBH_ROOT(rbt) && parent) {
 154                 if (RBE_LEFT(parent) == rbe) {
 155                         tmp = RBE_RIGHT(parent);
 156                         if (RBE_COLOR(tmp) == RB_RED) {
 157                                 rbe_set_blackred(tmp, parent);
 158                                 rbe_rotate_left(rbt, parent);
 159                                 tmp = RBE_RIGHT(parent);
 160                         }
 161                         if ((RBE_LEFT(tmp) == NULL
 162                              || RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK)
 163                             && (RBE_RIGHT(tmp) == NULL
 164                                 || RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK)) {
 165                                 RBE_COLOR(tmp) = RB_RED;
 166                                 rbe = parent;
 167                                 parent = RBE_PARENT(rbe);
 168                         } else {
 169                                 if (RBE_RIGHT(tmp) == NULL
 170                                     || RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK) {
 171                                         struct rb_entry *oleft;
 172
 173                                         oleft = RBE_LEFT(tmp);
 174                                         if (oleft != NULL)
 175                                                 RBE_COLOR(oleft) = RB_BLACK;
 176
 177                                         RBE_COLOR(tmp) = RB_RED;
 178                                         rbe_rotate_right(rbt, tmp);
 179                                         tmp = RBE_RIGHT(parent);
 180                                 }
 181
 182                                 RBE_COLOR(tmp) = RBE_COLOR(parent);
 183                                 RBE_COLOR(parent) = RB_BLACK;
 184                                 if (RBE_RIGHT(tmp))
 185                                         RBE_COLOR(RBE_RIGHT(tmp)) = RB_BLACK;
 186
 187                                 rbe_rotate_left(rbt, parent);
 188                                 rbe = RBH_ROOT(rbt);
 189                                 break;
 190                         }
 191                 } else {
 192                         tmp = RBE_LEFT(parent);
 193                         if (RBE_COLOR(tmp) == RB_RED) {
 194                                 rbe_set_blackred(tmp, parent);
 195                                 rbe_rotate_right(rbt, parent);
 196                                 tmp = RBE_LEFT(parent);
 197                         }
 198
 199                         if ((RBE_LEFT(tmp) == NULL
 200                              || RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK)
 201                             && (RBE_RIGHT(tmp) == NULL
 202                                 || RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK)) {
 203                                 RBE_COLOR(tmp) = RB_RED;
 204                                 rbe = parent;
 205                                 parent = RBE_PARENT(rbe);
 206                         } else {
 207                                 if (RBE_LEFT(tmp) == NULL
 208                                     || RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) {
 209                                         struct rb_entry *oright;
 210
 211                                         oright = RBE_RIGHT(tmp);
 212                                         if (oright != NULL)
 213                                                 RBE_COLOR(oright) = RB_BLACK;
 214
 215                                         RBE_COLOR(tmp) = RB_RED;
 216                                         rbe_rotate_left(rbt, tmp);
 217                                         tmp = RBE_LEFT(parent);
 218                                 }
 219
 220                                 RBE_COLOR(tmp) = RBE_COLOR(parent);
 221                                 RBE_COLOR(parent) = RB_BLACK;
 222                                 if (RBE_LEFT(tmp) != NULL)
 223                                         RBE_COLOR(RBE_LEFT(tmp)) = RB_BLACK;
 224
 225                                 rbe_rotate_right(rbt, parent);
 226                                 rbe = RBH_ROOT(rbt);
 227                                 break;
 228                         }
 229                 }
 230         }
 231
 232         if (rbe != NULL)
 233                 RBE_COLOR(rbe) = RB_BLACK;
 234 }
 235
 236 static inline struct rb_entry *
 237 rbe_remove(struct rbt_tree *rbt, struct rb_entry *rbe)
 238 {
 239         struct rb_entry *child, *parent, *old = rbe;
 240         unsigned int color;
 241
 242         if (RBE_LEFT(rbe) == NULL)
 243                 child = RBE_RIGHT(rbe);
 244         else if (RBE_RIGHT(rbe) == NULL)
 245                 child = RBE_LEFT(rbe);
 246         else {
 247                 struct rb_entry *tmp;
 248
 249                 rbe = RBE_RIGHT(rbe);
 250                 while ((tmp = RBE_LEFT(rbe)) != NULL)
 251                         rbe = tmp;
 252
 253                 child = RBE_RIGHT(rbe);
 254                 parent = RBE_PARENT(rbe);
 255                 color = RBE_COLOR(rbe);
 256                 if (child != NULL)
 257                         RBE_PARENT(child) = parent;
 258                 if (parent != NULL) {
 259                         if (RBE_LEFT(parent) == rbe)
 260                                 RBE_LEFT(parent) = child;
 261                         else
 262                                 RBE_RIGHT(parent) = child;
 263                 } else
 264                         RBH_ROOT(rbt) = child;
 265                 if (RBE_PARENT(rbe) == old)
 266                         parent = rbe;
 267                 *rbe = *old;
 268
 269                 tmp = RBE_PARENT(old);
 270                 if (tmp != NULL) {
 271                         if (RBE_LEFT(tmp) == old)
 272                                 RBE_LEFT(tmp) = rbe;
 273                         else
 274                                 RBE_RIGHT(tmp) = rbe;
 275                 } else
 276                         RBH_ROOT(rbt) = rbe;
 277
 278                 RBE_PARENT(RBE_LEFT(old)) = rbe;
 279                 if (RBE_RIGHT(old))
 280                         RBE_PARENT(RBE_RIGHT(old)) = rbe;
 281
 282                 goto color;
 283         }
 284
 285         parent = RBE_PARENT(rbe);
 286         color = RBE_COLOR(rbe);
 287
 288         if (child != NULL)
 289                 RBE_PARENT(child) = parent;
 290         if (parent != NULL) {
 291                 if (RBE_LEFT(parent) == rbe)
 292                         RBE_LEFT(parent) = child;
 293                 else
 294                         RBE_RIGHT(parent) = child;
 295         } else
 296                 RBH_ROOT(rbt) = child;
 297 color:
 298         if (color == RB_BLACK)
 299                 rbe_remove_color(rbt, parent, child);
 300
 301         rbt->count--;
 302         memset(old, 0, sizeof(*old));
 303         return (old);
 304 }
 305
 306 struct typed_rb_entry *typed_rb_remove(struct rbt_tree *rbt,
 307                                        struct rb_entry *rbe)
 308 {
 309         return rbe_remove(rbt, rbe);
 310 }
 311
 312 struct typed_rb_entry *typed_rb_insert(struct rbt_tree *rbt,
 313                 struct rb_entry *rbe, int (*cmpfn)(
 314                         const struct typed_rb_entry *a,
 315                         const struct typed_rb_entry *b))
 316 {
 317         struct rb_entry *tmp;
 318         struct rb_entry *parent = NULL;
 319         int comp = 0;
 320
 321         tmp = RBH_ROOT(rbt);
 322         while (tmp != NULL) {
 323                 parent = tmp;
 324
 325                 comp = cmpfn(rbe, tmp);
 326                 if (comp < 0)
 327                         tmp = RBE_LEFT(tmp);
 328                 else if (comp > 0)
 329                         tmp = RBE_RIGHT(tmp);
 330                 else
 331                         return tmp;
 332         }
 333
 334         rbe_set(rbe, parent);
 335
 336         if (parent != NULL) {
 337                 if (comp < 0)
 338                         RBE_LEFT(parent) = rbe;
 339                 else
 340                         RBE_RIGHT(parent) = rbe;
 341         } else
 342                 RBH_ROOT(rbt) = rbe;
 343
 344         rbe_insert_color(rbt, rbe);
 345
 346         return NULL;
 347 }
 348
 349 /* Finds the node with the same key as elm */
 350 const struct rb_entry *typed_rb_find(const struct rbt_tree *rbt,
 351                 const struct rb_entry *key,
 352                 int (*cmpfn)(
 353                         const struct typed_rb_entry *a,
 354                         const struct typed_rb_entry *b))
 355 {
 356         const struct rb_entry *tmp = RBH_ROOT(rbt);
 357         int comp;
 358
 359         while (tmp != NULL) {
 360                 comp = cmpfn(key, tmp);
 361                 if (comp < 0)
 362                         tmp = RBE_LEFT(tmp);
 363                 else if (comp > 0)
 364                         tmp = RBE_RIGHT(tmp);
 365                 else
 366                         return tmp;
 367         }
 368
 369         return NULL;
 370 }
 371
 372 const struct rb_entry *typed_rb_find_gteq(const struct rbt_tree *rbt,
 373                 const struct rb_entry *key,
 374                 int (*cmpfn)(
 375                         const struct typed_rb_entry *a,
 376                         const struct typed_rb_entry *b))
 377 {
 378         const struct rb_entry *tmp = RBH_ROOT(rbt), *best = NULL;
 379         int comp;
 380
 381         while (tmp != NULL) {
 382                 comp = cmpfn(key, tmp);
 383                 if (comp < 0) {
 384                         best = tmp;
 385                         tmp = RBE_LEFT(tmp);
 386                 } else if (comp > 0)
 387                         tmp = RBE_RIGHT(tmp);
 388                 else
 389                         return tmp;
 390         }
 391
 392         return best;
 393 }
 394
 395 const struct rb_entry *typed_rb_find_lt(const struct rbt_tree *rbt,
 396                 const struct rb_entry *key,
 397                 int (*cmpfn)(
 398                         const struct typed_rb_entry *a,
 399                         const struct typed_rb_entry *b))
 400 {
 401         const struct rb_entry *tmp = RBH_ROOT(rbt), *best = NULL;
 402         int comp;
 403
 404         while (tmp != NULL) {
 405                 comp = cmpfn(key, tmp);
 406                 if (comp <= 0)
 407                         tmp = RBE_LEFT(tmp);
 408                 else {
 409                         best = tmp;
 410                         tmp = RBE_RIGHT(tmp);
 411                 }
 412         }
 413
 414         return best;
 415 }
 416
 417 struct rb_entry *typed_rb_next(const struct rb_entry *rbe_const)
 418 {
 419         struct rb_entry *rbe = (struct rb_entry *)rbe_const;
 420
 421         if (RBE_RIGHT(rbe) != NULL) {
 422                 rbe = RBE_RIGHT(rbe);
 423                 while (RBE_LEFT(rbe) != NULL)
 424                         rbe = RBE_LEFT(rbe);
 425         } else {
 426                 if (RBE_PARENT(rbe) && (rbe == RBE_LEFT(RBE_PARENT(rbe))))
 427                         rbe = RBE_PARENT(rbe);
 428                 else {
 429                         while (RBE_PARENT(rbe)
 430                                && (rbe == RBE_RIGHT(RBE_PARENT(rbe))))
 431                                 rbe = RBE_PARENT(rbe);
 432                         rbe = RBE_PARENT(rbe);
 433                 }
 434         }
 435
 436         return rbe;
 437 }
 438
 439 struct rb_entry *typed_rb_prev(const struct rb_entry *rbe_const)
 440 {
 441         struct rb_entry *rbe = (struct rb_entry *)rbe_const;
 442
 443         if (RBE_LEFT(rbe)) {
 444                 rbe = RBE_LEFT(rbe);
 445                 while (RBE_RIGHT(rbe))
 446                         rbe = RBE_RIGHT(rbe);
 447         } else {
 448                 if (RBE_PARENT(rbe) && (rbe == RBE_RIGHT(RBE_PARENT(rbe))))
 449                         rbe = RBE_PARENT(rbe);
 450                 else {
 451                         while (RBE_PARENT(rbe)
 452                                && (rbe == RBE_LEFT(RBE_PARENT(rbe))))
 453                                 rbe = RBE_PARENT(rbe);
 454                         rbe = RBE_PARENT(rbe);
 455                 }
 456         }
 457
 458         return rbe;
 459 }
 460
 461 struct rb_entry *typed_rb_min(const struct rbt_tree *rbt)
 462 {
 463         struct rb_entry *rbe = RBH_ROOT(rbt);
 464         struct rb_entry *parent = NULL;
 465
 466         while (rbe != NULL) {
 467                 parent = rbe;
 468                 rbe = RBE_LEFT(rbe);
 469         }
 470
 471         return parent;
 472 }
 473
 474 struct rb_entry *typed_rb_max(const struct rbt_tree *rbt)
 475 {
 476         struct rb_entry *rbe = RBH_ROOT(rbt);
 477         struct rb_entry *parent = NULL;
 478
 479         while (rbe != NULL) {
 480                 parent = rbe;
 481                 rbe = RBE_RIGHT(rbe);
 482         }
 483
 484         return parent;
 485 }
 486
 487 bool typed_rb_member(const struct typed_rb_root *rbt,
 488                      const struct typed_rb_entry *rbe)
 489 {
 490         while (rbe->rbt_parent)
 491                 rbe = rbe->rbt_parent;
 492         return rbe == rbt->rbt_root;
 493 }