src/prop_rb.c

   1 /*      $NetBSD: prop_rb.c,v 1.9 2008/06/17 21:29:47 thorpej Exp $      */
   2
   3 /*-
   4  * Copyright (c) 2001 The NetBSD Foundation, Inc.
   5  * All rights reserved.
   6  *
   7  * This code is derived from software contributed to The NetBSD Foundation
   8  * by Matt Thomas <matt@3am-software.com>.
   9  *
  10  * Redistribution and use in source and binary forms, with or without
  11  * modification, are permitted provided that the following conditions
  12  * are met:
  13  * 1. Redistributions of source code must retain the above copyright
  14  *    notice, this list of conditions and the following disclaimer.
  15  * 2. Redistributions in binary form must reproduce the above copyright
  16  *    notice, this list of conditions and the following disclaimer in the
  17  *    documentation and/or other materials provided with the distribution.
  18  *
  19  * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
  20  * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
  21  * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  22  * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
  23  * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  24  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  25  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  26  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  27  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  28  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  29  * POSSIBILITY OF SUCH DAMAGE.
  30  */
  31
  32 #include "proplib.h"
  33 #include "prop_object_impl.h"
  34 #include "prop_rb_impl.h"
  35
  36 #undef KASSERT
  37 #ifdef RBDEBUG
  38 #define KASSERT(x)      _PROP_ASSERT(x)
  39 #else
  40 #define KASSERT(x)      /* nothing */
  41 #endif
  42
  43 #ifndef __predict_false
  44 #define __predict_false(x)      (x)
  45 #endif
  46
  47 static void rb_tree_reparent_nodes(struct rb_tree *, struct rb_node *,
  48                                    unsigned int);
  49 static void rb_tree_insert_rebalance(struct rb_tree *, struct rb_node *);
  50 static void rb_tree_removal_rebalance(struct rb_tree *, struct rb_node *,
  51         unsigned int);
  52 #ifdef RBDEBUG
  53 static const struct rb_node *rb_tree_iterate_const(const struct rb_tree *,
  54         const struct rb_node *, unsigned int);
  55 static bool rb_tree_check_node(const struct rb_tree *, const struct rb_node *,
  56         const struct rb_node *, bool);
  57 #endif
  58
  59 #ifdef RBDEBUG
  60 #define RBT_COUNT_INCR(rbt)     (rbt)->rbt_count++
  61 #define RBT_COUNT_DECR(rbt)     (rbt)->rbt_count--
  62 #else
  63 #define RBT_COUNT_INCR(rbt)     /* nothing */
  64 #define RBT_COUNT_DECR(rbt)     /* nothing */
  65 #endif
  66
  67 #define RBUNCONST(a)    ((void *)(unsigned long)(const void *)(a))
  68
  69 /*
  70  * Rather than testing for the NULL everywhere, all terminal leaves are
  71  * pointed to this node (and that includes itself).  Note that by setting
  72  * it to be const, that on some architectures trying to write to it will
  73  * cause a fault.
  74  */
  75 static const struct rb_node sentinel_node = {
  76         .rb_nodes = { RBUNCONST(&sentinel_node),
  77                       RBUNCONST(&sentinel_node),
  78                       NULL },
  79         .rb_u = { .u_s = { .s_sentinel = 1 } },
  80 };
  81
  82 void
  83 _prop_rb_tree_init(struct rb_tree *rbt, const struct rb_tree_ops *ops)
  84 {
  85         RB_TAILQ_INIT(&rbt->rbt_nodes);
  86 #ifdef RBDEBUG
  87         rbt->rbt_count = 0;
  88 #endif
  89         rbt->rbt_ops = ops;
  90         *((const struct rb_node **)&rbt->rbt_root) = &sentinel_node;
  91 }
  92
  93 /*
  94  * Swap the location and colors of 'self' and its child @ which.  The child
  95  * can not be a sentinel node.
  96  */
  97 /*ARGSUSED*/
  98 static void
  99 rb_tree_reparent_nodes(struct rb_tree *rbt _PROP_ARG_UNUSED,
 100     struct rb_node *old_father, unsigned int which)
 101 {
 102         const unsigned int other = which ^ RB_NODE_OTHER;
 103         struct rb_node * const grandpa = old_father->rb_parent;
 104         struct rb_node * const old_child = old_father->rb_nodes[which];
 105         struct rb_node * const new_father = old_child;
 106         struct rb_node * const new_child = old_father;
 107         unsigned int properties;
 108
 109         KASSERT(which == RB_NODE_LEFT || which == RB_NODE_RIGHT);
 110
 111         KASSERT(!RB_SENTINEL_P(old_child));
 112         KASSERT(old_child->rb_parent == old_father);
 113
 114         KASSERT(rb_tree_check_node(rbt, old_father, NULL, false));
 115         KASSERT(rb_tree_check_node(rbt, old_child, NULL, false));
 116         KASSERT(RB_ROOT_P(old_father) || rb_tree_check_node(rbt, grandpa, NULL, false));
 117
 118         /*
 119          * Exchange descendant linkages.
 120          */
 121         grandpa->rb_nodes[old_father->rb_position] = new_father;
 122         new_child->rb_nodes[which] = old_child->rb_nodes[other];
 123         new_father->rb_nodes[other] = new_child;
 124
 125         /*
 126          * Update ancestor linkages
 127          */
 128         new_father->rb_parent = grandpa;
 129         new_child->rb_parent = new_father;
 130
 131         /*
 132          * Exchange properties between new_father and new_child.  The only
 133          * change is that new_child's position is now on the other side.
 134          */
 135         properties = old_child->rb_properties;
 136         new_father->rb_properties = old_father->rb_properties;
 137         new_child->rb_properties = properties;
 138         new_child->rb_position = other;
 139
 140         /*
 141          * Make sure to reparent the new child to ourself.
 142          */
 143         if (!RB_SENTINEL_P(new_child->rb_nodes[which])) {
 144                 new_child->rb_nodes[which]->rb_parent = new_child;
 145                 new_child->rb_nodes[which]->rb_position = which;
 146         }
 147
 148         KASSERT(rb_tree_check_node(rbt, new_father, NULL, false));
 149         KASSERT(rb_tree_check_node(rbt, new_child, NULL, false));
 150         KASSERT(RB_ROOT_P(new_father) || rb_tree_check_node(rbt, grandpa, NULL, false));
 151 }
 152
 153 bool
 154 _prop_rb_tree_insert_node(struct rb_tree *rbt, struct rb_node *self)
 155 {
 156         struct rb_node *parent, *tmp;
 157         rb_compare_nodes_fn compare_nodes = rbt->rbt_ops->rbto_compare_nodes;
 158         unsigned int position;
 159
 160         self->rb_properties = 0;
 161         tmp = rbt->rbt_root;
 162         /*
 163          * This is a hack.  Because rbt->rbt_root is just a struct rb_node *,
 164          * just like rb_node->rb_nodes[RB_NODE_LEFT], we can use this fact to
 165          * avoid a lot of tests for root and know that even at root,
 166          * updating rb_node->rb_parent->rb_nodes[rb_node->rb_position] will
 167          * rbt->rbt_root.
 168          */
 169         /* LINTED: see above */
 170         parent = (struct rb_node *)&rbt->rbt_root;
 171         position = RB_NODE_LEFT;
 172
 173         /*
 174          * Find out where to place this new leaf.
 175          */
 176         while (!RB_SENTINEL_P(tmp)) {
 177                 const int diff = (*compare_nodes)(tmp, self);
 178                 if (__predict_false(diff == 0)) {
 179                         /*
 180                          * Node already exists; don't insert.
 181                          */
 182                         return false;
 183                 }
 184                 parent = tmp;
 185                 KASSERT(diff != 0);
 186                 if (diff < 0) {
 187                         position = RB_NODE_LEFT;
 188                 } else {
 189                         position = RB_NODE_RIGHT;
 190                 }
 191                 tmp = parent->rb_nodes[position];
 192         }
 193
 194 #ifdef RBDEBUG
 195         {
 196                 struct rb_node *prev = NULL, *next = NULL;
 197
 198                 if (position == RB_NODE_RIGHT)
 199                         prev = parent;
 200                 else if (tmp != rbt->rbt_root)
 201                         next = parent;
 202
 203                 /*
 204                  * Verify our sequential position
 205                  */
 206                 KASSERT(prev == NULL || !RB_SENTINEL_P(prev));
 207                 KASSERT(next == NULL || !RB_SENTINEL_P(next));
 208                 if (prev != NULL && next == NULL)
 209                         next = TAILQ_NEXT(prev, rb_link);
 210                 if (prev == NULL && next != NULL)
 211                         prev = TAILQ_PREV(next, rb_node_qh, rb_link);
 212                 KASSERT(prev == NULL || !RB_SENTINEL_P(prev));
 213                 KASSERT(next == NULL || !RB_SENTINEL_P(next));
 214                 KASSERT(prev == NULL
 215                         || (*compare_nodes)(prev, self) > 0);
 216                 KASSERT(next == NULL
 217                         || (*compare_nodes)(self, next) > 0);
 218         }
 219 #endif
 220
 221         /*
 222          * Initialize the node and insert as a leaf into the tree.
 223          */
 224         self->rb_parent = parent;
 225         self->rb_position = position;
 226         /* LINTED: rbt_root hack */
 227         if (__predict_false(parent == (struct rb_node *) &rbt->rbt_root)) {
 228                 RB_MARK_ROOT(self);
 229         } else {
 230                 KASSERT(position == RB_NODE_LEFT || position == RB_NODE_RIGHT);
 231                 KASSERT(!RB_ROOT_P(self));      /* Already done */
 232         }
 233         KASSERT(RB_SENTINEL_P(parent->rb_nodes[position]));
 234         self->rb_left = parent->rb_nodes[position];
 235         self->rb_right = parent->rb_nodes[position];
 236         parent->rb_nodes[position] = self;
 237         KASSERT(self->rb_left == &sentinel_node &&
 238             self->rb_right == &sentinel_node);
 239
 240         /*
 241          * Insert the new node into a sorted list for easy sequential access
 242          */
 243         RBT_COUNT_INCR(rbt);
 244 #ifdef RBDEBUG
 245         if (RB_ROOT_P(self)) {
 246                 RB_TAILQ_INSERT_HEAD(&rbt->rbt_nodes, self, rb_link);
 247         } else if (position == RB_NODE_LEFT) {
 248                 KASSERT((*compare_nodes)(self, self->rb_parent) > 0);
 249                 RB_TAILQ_INSERT_BEFORE(self->rb_parent, self, rb_link);
 250         } else {
 251                 KASSERT((*compare_nodes)(self->rb_parent, self) > 0);
 252                 RB_TAILQ_INSERT_AFTER(&rbt->rbt_nodes, self->rb_parent,
 253                     self, rb_link);
 254         }
 255 #endif
 256
 257 #if 0
 258         /*
 259          * Validate the tree before we rebalance
 260          */
 261         _prop_rb_tree_check(rbt, false);
 262 #endif
 263
 264         /*
 265          * Rebalance tree after insertion
 266          */
 267         rb_tree_insert_rebalance(rbt, self);
 268
 269 #if 0
 270         /*
 271          * Validate the tree after we rebalanced
 272          */
 273         _prop_rb_tree_check(rbt, true);
 274 #endif
 275
 276         return true;
 277 }
 278 \f
 279 static void
 280 rb_tree_insert_rebalance(struct rb_tree *rbt, struct rb_node *self)
 281 {
 282         RB_MARK_RED(self);
 283
 284         while (!RB_ROOT_P(self) && RB_RED_P(self->rb_parent)) {
 285                 const unsigned int which =
 286                      (self->rb_parent == self->rb_parent->rb_parent->rb_left
 287                         ? RB_NODE_LEFT
 288                         : RB_NODE_RIGHT);
 289                 const unsigned int other = which ^ RB_NODE_OTHER;
 290                 struct rb_node * father = self->rb_parent;
 291                 struct rb_node * grandpa = father->rb_parent;
 292                 struct rb_node * const uncle = grandpa->rb_nodes[other];
 293
 294                 KASSERT(!RB_SENTINEL_P(self));
 295                 /*
 296                  * We are red and our parent is red, therefore we must have a
 297                  * grandfather and he must be black.
 298                  */
 299                 KASSERT(RB_RED_P(self)
 300                         && RB_RED_P(father)
 301                         && RB_BLACK_P(grandpa));
 302
 303                 if (RB_RED_P(uncle)) {
 304                         /*
 305                          * Case 1: our uncle is red
 306                          *   Simply invert the colors of our parent and
 307                          *   uncle and make our grandparent red.  And
 308                          *   then solve the problem up at his level.
 309                          */
 310                         RB_MARK_BLACK(uncle);
 311                         RB_MARK_BLACK(father);
 312                         RB_MARK_RED(grandpa);
 313                         self = grandpa;
 314                         continue;
 315                 }
 316                 /*
 317                  * Case 2&3: our uncle is black.
 318                  */
 319                 if (self == father->rb_nodes[other]) {
 320                         /*
 321                          * Case 2: we are on the same side as our uncle
 322                          *   Swap ourselves with our parent so this case
 323                          *   becomes case 3.  Basically our parent becomes our
 324                          *   child.
 325                          */
 326                         rb_tree_reparent_nodes(rbt, father, other);
 327                         KASSERT(father->rb_parent == self);
 328                         KASSERT(self->rb_nodes[which] == father);
 329                         KASSERT(self->rb_parent == grandpa);
 330                         self = father;
 331                         father = self->rb_parent;
 332                 }
 333                 KASSERT(RB_RED_P(self) && RB_RED_P(father));
 334                 KASSERT(grandpa->rb_nodes[which] == father);
 335                 /*
 336                  * Case 3: we are opposite a child of a black uncle.
 337                  *   Swap our parent and grandparent.  Since our grandfather
 338                  *   is black, our father will become black and our new sibling
 339                  *   (former grandparent) will become red.
 340                  */
 341                 rb_tree_reparent_nodes(rbt, grandpa, which);
 342                 KASSERT(self->rb_parent == father);
 343                 KASSERT(self->rb_parent->rb_nodes[self->rb_position ^ RB_NODE_OTHER] == grandpa);
 344                 KASSERT(RB_RED_P(self));
 345                 KASSERT(RB_BLACK_P(father));
 346                 KASSERT(RB_RED_P(grandpa));
 347                 break;
 348         }
 349
 350         /*
 351          * Final step: Set the root to black.
 352          */
 353         RB_MARK_BLACK(rbt->rbt_root);
 354 }
 355 \f
 356 struct rb_node *
 357 _prop_rb_tree_find(struct rb_tree *rbt, const void *key)
 358 {
 359         struct rb_node *parent = rbt->rbt_root;
 360         rb_compare_key_fn compare_key = rbt->rbt_ops->rbto_compare_key;
 361
 362         while (!RB_SENTINEL_P(parent)) {
 363                 const int diff = (*compare_key)(parent, key);
 364                 if (diff == 0)
 365                         return parent;
 366                 parent = parent->rb_nodes[diff > 0];
 367         }
 368
 369         return NULL;
 370 }
 371 \f
 372 static void
 373 rb_tree_prune_node(struct rb_tree *rbt, struct rb_node *self, int rebalance)
 374 {
 375         const unsigned int which = self->rb_position;
 376         struct rb_node *father = self->rb_parent;
 377
 378         KASSERT(rebalance || (RB_ROOT_P(self) || RB_RED_P(self)));
 379         KASSERT(!rebalance || RB_BLACK_P(self));
 380         KASSERT(RB_CHILDLESS_P(self));
 381         KASSERT(rb_tree_check_node(rbt, self, NULL, false));
 382
 383         father->rb_nodes[which] = self->rb_left;
 384
 385         /*
 386          * Remove ourselves from the node list and decrement the count.
 387          */
 388         RB_TAILQ_REMOVE(&rbt->rbt_nodes, self, rb_link);
 389         RBT_COUNT_DECR(rbt);
 390
 391         if (rebalance)
 392                 rb_tree_removal_rebalance(rbt, father, which);
 393         KASSERT(RB_ROOT_P(self) || rb_tree_check_node(rbt, father, NULL, true));
 394 }
 395
 396 static void
 397 rb_tree_swap_prune_and_rebalance(struct rb_tree *rbt, struct rb_node *self,
 398         struct rb_node *standin)
 399 {
 400         unsigned int standin_which = standin->rb_position;
 401         unsigned int standin_other = standin_which ^ RB_NODE_OTHER;
 402         struct rb_node *standin_child;
 403         struct rb_node *standin_father;
 404         bool rebalance = RB_BLACK_P(standin);
 405
 406         if (standin->rb_parent == self) {
 407                 /*
 408                  * As a child of self, any childen would be opposite of
 409                  * our parent (self).
 410                  */
 411                 KASSERT(RB_SENTINEL_P(standin->rb_nodes[standin_other]));
 412                 standin_child = standin->rb_nodes[standin_which];
 413         } else {
 414                 /*
 415                  * Since we aren't a child of self, any childen would be
 416                  * on the same side as our parent (self).
 417                  */
 418                 KASSERT(RB_SENTINEL_P(standin->rb_nodes[standin_which]));
 419                 standin_child = standin->rb_nodes[standin_other];
 420         }
 421
 422         /*
 423          * the node we are removing must have two children.
 424          */
 425         KASSERT(RB_TWOCHILDREN_P(self));
 426         /*
 427          * If standin has a child, it must be red.
 428          */
 429         KASSERT(RB_SENTINEL_P(standin_child) || RB_RED_P(standin_child));
 430
 431         /*
 432          * Verify things are sane.
 433          */
 434         KASSERT(rb_tree_check_node(rbt, self, NULL, false));
 435         KASSERT(rb_tree_check_node(rbt, standin, NULL, false));
 436
 437         if (!RB_SENTINEL_P(standin_child)) {
 438                 /*
 439                  * We know we have a red child so if we swap them we can
 440                  * void flipping standin's child to black afterwards.
 441                  */
 442                 KASSERT(rb_tree_check_node(rbt, standin_child, NULL, true));
 443                 rb_tree_reparent_nodes(rbt, standin,
 444                     standin_child->rb_position);
 445                 KASSERT(rb_tree_check_node(rbt, standin, NULL, true));
 446                 KASSERT(rb_tree_check_node(rbt, standin_child, NULL, true));
 447                 /*
 448                  * Since we are removing a red leaf, no need to rebalance.
 449                  */
 450                 rebalance = false;
 451                 /*
 452                  * We know that standin can not be a child of self, so
 453                  * update before of that.
 454                  */
 455                 KASSERT(standin->rb_parent != self);
 456                 standin_which = standin->rb_position;
 457                 standin_other = standin_which ^ RB_NODE_OTHER;
 458         }
 459         KASSERT(RB_CHILDLESS_P(standin));
 460
 461         /*
 462          * If we are about to delete the standin's father, then when we call
 463          * rebalance, we need to use ourselves as our father.  Otherwise
 464          * remember our original father.  Also, if we are our standin's father
 465          * we only need to reparent the standin's brother.
 466          */
 467         if (standin->rb_parent == self) {
 468                 /*
 469                  * |   R   -->   S   |
 470                  * | Q   S --> Q   * |
 471                  * |       -->       |
 472                  */
 473                 standin_father = standin;
 474                 KASSERT(RB_SENTINEL_P(standin->rb_nodes[standin_other]));
 475                 KASSERT(!RB_SENTINEL_P(self->rb_nodes[standin_other]));
 476                 KASSERT(self->rb_nodes[standin_which] == standin);
 477                 /*
 478                  * Make our brother our son.
 479                  */
 480                 standin->rb_nodes[standin_other] = self->rb_nodes[standin_other];
 481                 standin->rb_nodes[standin_other]->rb_parent = standin;
 482                 KASSERT(standin->rb_nodes[standin_other]->rb_position == standin_other);
 483         } else {
 484                 /*
 485                  * |  P      -->  P    |
 486                  * |      S  -->    Q  |
 487                  * |    Q    -->       |
 488                  */
 489                 standin_father = standin->rb_parent;
 490                 standin_father->rb_nodes[standin_which] =
 491                     standin->rb_nodes[standin_which];
 492                 standin->rb_left = self->rb_left;
 493                 standin->rb_right = self->rb_right;
 494                 standin->rb_left->rb_parent = standin;
 495                 standin->rb_right->rb_parent = standin;
 496         }
 497
 498         /*
 499          * Now copy the result of self to standin and then replace
 500          * self with standin in the tree.
 501          */
 502         standin->rb_parent = self->rb_parent;
 503         standin->rb_properties = self->rb_properties;
 504         standin->rb_parent->rb_nodes[standin->rb_position] = standin;
 505
 506         /*
 507          * Remove ourselves from the node list and decrement the count.
 508          */
 509         RB_TAILQ_REMOVE(&rbt->rbt_nodes, self, rb_link);
 510         RBT_COUNT_DECR(rbt);
 511
 512         KASSERT(rb_tree_check_node(rbt, standin, NULL, false));
 513         KASSERT(rb_tree_check_node(rbt, standin_father, NULL, false));
 514
 515         if (!rebalance)
 516                 return;
 517
 518         rb_tree_removal_rebalance(rbt, standin_father, standin_which);
 519         KASSERT(rb_tree_check_node(rbt, standin, NULL, true));
 520 }
 521
 522 /*
 523  * We could do this by doing
 524  *      rb_tree_node_swap(rbt, self, which);
 525  *      rb_tree_prune_node(rbt, self, false);
 526  *
 527  * But it's more efficient to just evalate and recolor the child.
 528  */
 529 /*ARGSUSED*/
 530 static void
 531 rb_tree_prune_blackred_branch(struct rb_tree *rbt _PROP_ARG_UNUSED,
 532     struct rb_node *self, unsigned int which)
 533 {
 534         struct rb_node *parent = self->rb_parent;
 535         struct rb_node *child = self->rb_nodes[which];
 536
 537         KASSERT(which == RB_NODE_LEFT || which == RB_NODE_RIGHT);
 538         KASSERT(RB_BLACK_P(self) && RB_RED_P(child));
 539         KASSERT(!RB_TWOCHILDREN_P(child));
 540         KASSERT(RB_CHILDLESS_P(child));
 541         KASSERT(rb_tree_check_node(rbt, self, NULL, false));
 542         KASSERT(rb_tree_check_node(rbt, child, NULL, false));
 543
 544         /*
 545          * Remove ourselves from the tree and give our former child our
 546          * properties (position, color, root).
 547          */
 548         parent->rb_nodes[self->rb_position] = child;
 549         child->rb_parent = parent;
 550         child->rb_properties = self->rb_properties;
 551
 552         /*
 553          * Remove ourselves from the node list and decrement the count.
 554          */
 555         RB_TAILQ_REMOVE(&rbt->rbt_nodes, self, rb_link);
 556         RBT_COUNT_DECR(rbt);
 557
 558         KASSERT(RB_ROOT_P(self) || rb_tree_check_node(rbt, parent, NULL, true));
 559         KASSERT(rb_tree_check_node(rbt, child, NULL, true));
 560 }
 561 /*
 562  *
 563  */
 564 void
 565 _prop_rb_tree_remove_node(struct rb_tree *rbt, struct rb_node *self)
 566 {
 567         struct rb_node *standin;
 568         unsigned int which;
 569         /*
 570          * In the following diagrams, we (the node to be removed) are S.  Red
 571          * nodes are lowercase.  T could be either red or black.
 572          *
 573          * Remember the major axiom of the red-black tree: the number of
 574          * black nodes from the root to each leaf is constant across all
 575          * leaves, only the number of red nodes varies.
 576          *
 577          * Thus removing a red leaf doesn't require any other changes to a
 578          * red-black tree.  So if we must remove a node, attempt to rearrange
 579          * the tree so we can remove a red node.
 580          *
 581          * The simpliest case is a childless red node or a childless root node:
 582          *
 583          * |    T  -->    T  |    or    |  R  -->  *  |
 584          * |  s    -->  *    |
 585          */
 586         if (RB_CHILDLESS_P(self)) {
 587                 if (RB_RED_P(self) || RB_ROOT_P(self)) {
 588                         rb_tree_prune_node(rbt, self, false);
 589                         return;
 590                 }
 591                 rb_tree_prune_node(rbt, self, true);
 592                 return;
 593         }
 594         KASSERT(!RB_CHILDLESS_P(self));
 595         if (!RB_TWOCHILDREN_P(self)) {
 596                 /*
 597                  * The next simpliest case is the node we are deleting is
 598                  * black and has one red child.
 599                  *
 600                  * |      T  -->      T  -->      T  |
 601                  * |    S    -->  R      -->  R      |
 602                  * |  r      -->    s    -->    *    |
 603                  */
 604                 which = RB_LEFT_SENTINEL_P(self) ? RB_NODE_RIGHT : RB_NODE_LEFT;
 605                 KASSERT(RB_BLACK_P(self));
 606                 KASSERT(RB_RED_P(self->rb_nodes[which]));
 607                 KASSERT(RB_CHILDLESS_P(self->rb_nodes[which]));
 608                 rb_tree_prune_blackred_branch(rbt, self, which);
 609                 return;
 610         }
 611         KASSERT(RB_TWOCHILDREN_P(self));
 612
 613         /*
 614          * We invert these because we prefer to remove from the inside of
 615          * the tree.
 616          */
 617         which = self->rb_position ^ RB_NODE_OTHER;
 618
 619         /*
 620          * Let's find the node closes to us opposite of our parent
 621          * Now swap it with ourself, "prune" it, and rebalance, if needed.
 622          */
 623         standin = _prop_rb_tree_iterate(rbt, self, which);
 624         rb_tree_swap_prune_and_rebalance(rbt, self, standin);
 625 }
 626
 627 static void
 628 rb_tree_removal_rebalance(struct rb_tree *rbt, struct rb_node *parent,
 629         unsigned int which)
 630 {
 631         KASSERT(!RB_SENTINEL_P(parent));
 632         KASSERT(RB_SENTINEL_P(parent->rb_nodes[which]));
 633         KASSERT(which == RB_NODE_LEFT || which == RB_NODE_RIGHT);
 634
 635         while (RB_BLACK_P(parent->rb_nodes[which])) {
 636                 unsigned int other = which ^ RB_NODE_OTHER;
 637                 struct rb_node *brother = parent->rb_nodes[other];
 638
 639                 KASSERT(!RB_SENTINEL_P(brother));
 640                 /*
 641                  * For cases 1, 2a, and 2b, our brother's children must
 642                  * be black and our father must be black
 643                  */
 644                 if (RB_BLACK_P(parent)
 645                     && RB_BLACK_P(brother->rb_left)
 646                     && RB_BLACK_P(brother->rb_right)) {
 647                         /*
 648                          * Case 1: Our brother is red, swap its position
 649                          * (and colors) with our parent.  This is now case 2b.
 650                          *
 651                          *    B         ->        D
 652                          *  x     d     ->    b     E
 653                          *      C   E   ->  x   C
 654                          */
 655                         if (RB_RED_P(brother)) {
 656                                 KASSERT(RB_BLACK_P(parent));
 657                                 rb_tree_reparent_nodes(rbt, parent, other);
 658                                 brother = parent->rb_nodes[other];
 659                                 KASSERT(!RB_SENTINEL_P(brother));
 660                                 KASSERT(RB_BLACK_P(brother));
 661                                 KASSERT(RB_RED_P(parent));
 662                                 KASSERT(rb_tree_check_node(rbt, brother, NULL, false));
 663                                 KASSERT(rb_tree_check_node(rbt, parent, NULL, false));
 664                         } else {
 665                                 /*
 666                                  * Both our parent and brother are black.
 667                                  * Change our brother to red, advance up rank
 668                                  * and go through the loop again.
 669                                  *
 670                                  *    B         ->    B
 671                                  *  A     D     ->  A     d
 672                                  *      C   E   ->      C   E
 673                                  */
 674                                 RB_MARK_RED(brother);
 675                                 KASSERT(RB_BLACK_P(brother->rb_left));
 676                                 KASSERT(RB_BLACK_P(brother->rb_right));
 677                                 if (RB_ROOT_P(parent))
 678                                         return;
 679                                 KASSERT(rb_tree_check_node(rbt, brother, NULL, false));
 680                                 KASSERT(rb_tree_check_node(rbt, parent, NULL, false));
 681                                 which = parent->rb_position;
 682                                 parent = parent->rb_parent;
 683                         }
 684                 } else if (RB_RED_P(parent)
 685                     && RB_BLACK_P(brother)
 686                     && RB_BLACK_P(brother->rb_left)
 687                     && RB_BLACK_P(brother->rb_right)) {
 688                         KASSERT(RB_BLACK_P(brother));
 689                         KASSERT(RB_BLACK_P(brother->rb_left));
 690                         KASSERT(RB_BLACK_P(brother->rb_right));
 691                         RB_MARK_BLACK(parent);
 692                         RB_MARK_RED(brother);
 693                         KASSERT(rb_tree_check_node(rbt, brother, NULL, true));
 694                         break;          /* We're done! */
 695                 } else {
 696                         KASSERT(RB_BLACK_P(brother));
 697                         KASSERT(!RB_CHILDLESS_P(brother));
 698                         /*
 699                          * Case 3: our brother is black, our left nephew is
 700                          * red, and our right nephew is black.  Swap our
 701                          * brother with our left nephew.   This result in a
 702                          * tree that matches case 4.
 703                          *
 704                          *     B         ->       D
 705                          * A       D     ->   B     E
 706                          *       c   e   -> A   C
 707                          */
 708                         if (RB_BLACK_P(brother->rb_nodes[other])) {
 709                                 KASSERT(RB_RED_P(brother->rb_nodes[which]));
 710                                 rb_tree_reparent_nodes(rbt, brother, which);
 711                                 KASSERT(brother->rb_parent == parent->rb_nodes[other]);
 712                                 brother = parent->rb_nodes[other];
 713                                 KASSERT(RB_RED_P(brother->rb_nodes[other]));
 714                         }
 715                         /*
 716                          * Case 4: our brother is black and our right nephew
 717                          * is red.  Swap our parent and brother locations and
 718                          * change our right nephew to black.  (these can be
 719                          * done in either order so we change the color first).
 720                          * The result is a valid red-black tree and is a
 721                          * terminal case.
 722                          *
 723                          *     B         ->       D
 724                          * A       D     ->   B     E
 725                          *       c   e   -> A   C
 726                          */
 727                         RB_MARK_BLACK(brother->rb_nodes[other]);
 728                         rb_tree_reparent_nodes(rbt, parent, other);
 729                         break;          /* We're done! */
 730                 }
 731         }
 732         KASSERT(rb_tree_check_node(rbt, parent, NULL, true));
 733 }
 734
 735 struct rb_node *
 736 _prop_rb_tree_iterate(struct rb_tree *rbt, struct rb_node *self,
 737         unsigned int direction)
 738 {
 739         const unsigned int other = direction ^ RB_NODE_OTHER;
 740         KASSERT(direction == RB_NODE_LEFT || direction == RB_NODE_RIGHT);
 741
 742         if (self == NULL) {
 743                 self = rbt->rbt_root;
 744                 if (RB_SENTINEL_P(self))
 745                         return NULL;
 746                 while (!RB_SENTINEL_P(self->rb_nodes[other]))
 747                         self = self->rb_nodes[other];
 748                 return self;
 749         }
 750         KASSERT(!RB_SENTINEL_P(self));
 751         /*
 752          * We can't go any further in this direction.  We proceed up in the
 753          * opposite direction until our parent is in direction we want to go.
 754          */
 755         if (RB_SENTINEL_P(self->rb_nodes[direction])) {
 756                 while (!RB_ROOT_P(self)) {
 757                         if (other == self->rb_position)
 758                                 return self->rb_parent;
 759                         self = self->rb_parent;
 760                 }
 761                 return NULL;
 762         }
 763
 764         /*
 765          * Advance down one in current direction and go down as far as possible
 766          * in the opposite direction.
 767          */
 768         self = self->rb_nodes[direction];
 769         KASSERT(!RB_SENTINEL_P(self));
 770         while (!RB_SENTINEL_P(self->rb_nodes[other]))
 771                 self = self->rb_nodes[other];
 772         return self;
 773 }
 774
 775 #ifdef RBDEBUG
 776 static const struct rb_node *
 777 rb_tree_iterate_const(const struct rb_tree *rbt, const struct rb_node *self,
 778         unsigned int direction)
 779 {
 780         const unsigned int other = direction ^ RB_NODE_OTHER;
 781         KASSERT(direction == RB_NODE_LEFT || direction == RB_NODE_RIGHT);
 782
 783         if (self == NULL) {
 784                 self = rbt->rbt_root;
 785                 if (RB_SENTINEL_P(self))
 786                         return NULL;
 787                 while (!RB_SENTINEL_P(self->rb_nodes[other]))
 788                         self = self->rb_nodes[other];
 789                 return self;
 790         }
 791         KASSERT(!RB_SENTINEL_P(self));
 792         /*
 793          * We can't go any further in this direction.  We proceed up in the
 794          * opposite direction until our parent is in direction we want to go.
 795          */
 796         if (RB_SENTINEL_P(self->rb_nodes[direction])) {
 797                 while (!RB_ROOT_P(self)) {
 798                         if (other == self->rb_position)
 799                                 return self->rb_parent;
 800                         self = self->rb_parent;
 801                 }
 802                 return NULL;
 803         }
 804
 805         /*
 806          * Advance down one in current direction and go down as far as possible
 807          * in the opposite direction.
 808          */
 809         self = self->rb_nodes[direction];
 810         KASSERT(!RB_SENTINEL_P(self));
 811         while (!RB_SENTINEL_P(self->rb_nodes[other]))
 812                 self = self->rb_nodes[other];
 813         return self;
 814 }
 815
 816 static bool
 817 rb_tree_check_node(const struct rb_tree *rbt, const struct rb_node *self,
 818         const struct rb_node *prev, bool red_check)
 819 {
 820         KASSERT(!self->rb_sentinel);
 821         KASSERT(self->rb_left);
 822         KASSERT(self->rb_right);
 823         KASSERT(prev == NULL ||
 824                 (*rbt->rbt_ops->rbto_compare_nodes)(prev, self) > 0);
 825
 826         /*
 827          * Verify our relationship to our parent.
 828          */
 829         if (RB_ROOT_P(self)) {
 830                 KASSERT(self == rbt->rbt_root);
 831                 KASSERT(self->rb_position == RB_NODE_LEFT);
 832                 KASSERT(self->rb_parent->rb_nodes[RB_NODE_LEFT] == self);
 833                 KASSERT(self->rb_parent == (const struct rb_node *) &rbt->rbt_root);
 834         } else {
 835                 KASSERT(self != rbt->rbt_root);
 836                 KASSERT(!RB_PARENT_SENTINEL_P(self));
 837                 if (self->rb_position == RB_NODE_LEFT) {
 838                         KASSERT((*rbt->rbt_ops->rbto_compare_nodes)(self, self->rb_parent) > 0);
 839                         KASSERT(self->rb_parent->rb_nodes[RB_NODE_LEFT] == self);
 840                 } else {
 841                         KASSERT((*rbt->rbt_ops->rbto_compare_nodes)(self, self->rb_parent) < 0);
 842                         KASSERT(self->rb_parent->rb_nodes[RB_NODE_RIGHT] == self);
 843                 }
 844         }
 845
 846         /*
 847          * Verify our position in the linked list against the tree itself.
 848          */
 849         {
 850                 const struct rb_node *prev0 = rb_tree_iterate_const(rbt, self, RB_NODE_LEFT);
 851                 const struct rb_node *next0 = rb_tree_iterate_const(rbt, self, RB_NODE_RIGHT);
 852                 KASSERT(prev0 == TAILQ_PREV(self, rb_node_qh, rb_link));
 853                 if (next0 != TAILQ_NEXT(self, rb_link))
 854                         next0 = rb_tree_iterate_const(rbt, self, RB_NODE_RIGHT);
 855                 KASSERT(next0 == TAILQ_NEXT(self, rb_link));
 856         }
 857
 858         /*
 859          * The root must be black.
 860          * There can never be two adjacent red nodes.
 861          */
 862         if (red_check) {
 863                 KASSERT(!RB_ROOT_P(self) || RB_BLACK_P(self));
 864                 if (RB_RED_P(self)) {
 865                         const struct rb_node *brother;
 866                         KASSERT(!RB_ROOT_P(self));
 867                         brother = self->rb_parent->rb_nodes[self->rb_position ^ RB_NODE_OTHER];
 868                         KASSERT(RB_BLACK_P(self->rb_parent));
 869                         /*
 870                          * I'm red and have no children, then I must either
 871                          * have no brother or my brother also be red and
 872                          * also have no children.  (black count == 0)
 873                          */
 874                         KASSERT(!RB_CHILDLESS_P(self)
 875                                 || RB_SENTINEL_P(brother)
 876                                 || RB_RED_P(brother)
 877                                 || RB_CHILDLESS_P(brother));
 878                         /*
 879                          * If I'm not childless, I must have two children
 880                          * and they must be both be black.
 881                          */
 882                         KASSERT(RB_CHILDLESS_P(self)
 883                                 || (RB_TWOCHILDREN_P(self)
 884                                     && RB_BLACK_P(self->rb_left)
 885                                     && RB_BLACK_P(self->rb_right)));
 886                         /*
 887                          * If I'm not childless, thus I have black children,
 888                          * then my brother must either be black or have two
 889                          * black children.
 890                          */
 891                         KASSERT(RB_CHILDLESS_P(self)
 892                                 || RB_BLACK_P(brother)
 893                                 || (RB_TWOCHILDREN_P(brother)
 894                                     && RB_BLACK_P(brother->rb_left)
 895                                     && RB_BLACK_P(brother->rb_right)));
 896                 } else {
 897                         /*
 898                          * If I'm black and have one child, that child must
 899                          * be red and childless.
 900                          */
 901                         KASSERT(RB_CHILDLESS_P(self)
 902                                 || RB_TWOCHILDREN_P(self)
 903                                 || (!RB_LEFT_SENTINEL_P(self)
 904                                     && RB_RIGHT_SENTINEL_P(self)
 905                                     && RB_RED_P(self->rb_left)
 906                                     && RB_CHILDLESS_P(self->rb_left))
 907                                 || (!RB_RIGHT_SENTINEL_P(self)
 908                                     && RB_LEFT_SENTINEL_P(self)
 909                                     && RB_RED_P(self->rb_right)
 910                                     && RB_CHILDLESS_P(self->rb_right)));
 911
 912                         /*
 913                          * If I'm a childless black node and my parent is
 914                          * black, my 2nd closet relative away from my parent
 915                          * is either red or has a red parent or red children.
 916                          */
 917                         if (!RB_ROOT_P(self)
 918                             && RB_CHILDLESS_P(self)
 919                             && RB_BLACK_P(self->rb_parent)) {
 920                                 const unsigned int which = self->rb_position;
 921                                 const unsigned int other = which ^ RB_NODE_OTHER;
 922                                 const struct rb_node *relative0, *relative;
 923
 924                                 relative0 = rb_tree_iterate_const(rbt,
 925                                     self, other);
 926                                 KASSERT(relative0 != NULL);
 927                                 relative = rb_tree_iterate_const(rbt,
 928                                     relative0, other);
 929                                 KASSERT(relative != NULL);
 930                                 KASSERT(RB_SENTINEL_P(relative->rb_nodes[which]));
 931 #if 0
 932                                 KASSERT(RB_RED_P(relative)
 933                                         || RB_RED_P(relative->rb_left)
 934                                         || RB_RED_P(relative->rb_right)
 935                                         || RB_RED_P(relative->rb_parent));
 936 #endif
 937                         }
 938                 }
 939                 /*
 940                  * A grandparent's children must be real nodes and not
 941                  * sentinels.  First check out grandparent.
 942                  */
 943                 KASSERT(RB_ROOT_P(self)
 944                         || RB_ROOT_P(self->rb_parent)
 945                         || RB_TWOCHILDREN_P(self->rb_parent->rb_parent));
 946                 /*
 947                  * If we are have grandchildren on our left, then
 948                  * we must have a child on our right.
 949                  */
 950                 KASSERT(RB_LEFT_SENTINEL_P(self)
 951                         || RB_CHILDLESS_P(self->rb_left)
 952                         || !RB_RIGHT_SENTINEL_P(self));
 953                 /*
 954                  * If we are have grandchildren on our right, then
 955                  * we must have a child on our left.
 956                  */
 957                 KASSERT(RB_RIGHT_SENTINEL_P(self)
 958                         || RB_CHILDLESS_P(self->rb_right)
 959                         || !RB_LEFT_SENTINEL_P(self));
 960
 961                 /*
 962                  * If we have a child on the left and it doesn't have two
 963                  * children make sure we don't have great-great-grandchildren on
 964                  * the right.
 965                  */
 966                 KASSERT(RB_TWOCHILDREN_P(self->rb_left)
 967                         || RB_CHILDLESS_P(self->rb_right)
 968                         || RB_CHILDLESS_P(self->rb_right->rb_left)
 969                         || RB_CHILDLESS_P(self->rb_right->rb_left->rb_left)
 970                         || RB_CHILDLESS_P(self->rb_right->rb_left->rb_right)
 971                         || RB_CHILDLESS_P(self->rb_right->rb_right)
 972                         || RB_CHILDLESS_P(self->rb_right->rb_right->rb_left)
 973                         || RB_CHILDLESS_P(self->rb_right->rb_right->rb_right));
 974
 975                 /*
 976                  * If we have a child on the right and it doesn't have two
 977                  * children make sure we don't have great-great-grandchildren on
 978                  * the left.
 979                  */
 980                 KASSERT(RB_TWOCHILDREN_P(self->rb_right)
 981                         || RB_CHILDLESS_P(self->rb_left)
 982                         || RB_CHILDLESS_P(self->rb_left->rb_left)
 983                         || RB_CHILDLESS_P(self->rb_left->rb_left->rb_left)
 984                         || RB_CHILDLESS_P(self->rb_left->rb_left->rb_right)
 985                         || RB_CHILDLESS_P(self->rb_left->rb_right)
 986                         || RB_CHILDLESS_P(self->rb_left->rb_right->rb_left)
 987                         || RB_CHILDLESS_P(self->rb_left->rb_right->rb_right));
 988
 989                 /*
 990                  * If we are fully interior node, then our predecessors and
 991                  * successors must have no children in our direction.
 992                  */
 993                 if (RB_TWOCHILDREN_P(self)) {
 994                         const struct rb_node *prev0;
 995                         const struct rb_node *next0;
 996
 997                         prev0 = rb_tree_iterate_const(rbt, self, RB_NODE_LEFT);
 998                         KASSERT(prev0 != NULL);
 999                         KASSERT(RB_RIGHT_SENTINEL_P(prev0));
1000
1001                         next0 = rb_tree_iterate_const(rbt, self, RB_NODE_RIGHT);
1002                         KASSERT(next0 != NULL);
1003                         KASSERT(RB_LEFT_SENTINEL_P(next0));
1004                 }
1005         }
1006
1007         return true;
1008 }
1009
1010 static unsigned int
1011 rb_tree_count_black(const struct rb_node *self)
1012 {
1013         unsigned int left, right;
1014
1015         if (RB_SENTINEL_P(self))
1016                 return 0;
1017
1018         left = rb_tree_count_black(self->rb_left);
1019         right = rb_tree_count_black(self->rb_right);
1020
1021         KASSERT(left == right);
1022
1023         return left + RB_BLACK_P(self);
1024 }
1025
1026 void
1027 _prop_rb_tree_check(const struct rb_tree *rbt, bool red_check)
1028 {
1029         const struct rb_node *self;
1030         const struct rb_node *prev;
1031         unsigned int count;
1032
1033         KASSERT(rbt->rbt_root == NULL || rbt->rbt_root->rb_position == RB_NODE_LEFT);
1034
1035         prev = NULL;
1036         count = 0;
1037         TAILQ_FOREACH(self, &rbt->rbt_nodes, rb_link) {
1038                 rb_tree_check_node(rbt, self, prev, false);
1039                 count++;
1040         }
1041         KASSERT(rbt->rbt_count == count);
1042         KASSERT(RB_SENTINEL_P(rbt->rbt_root)
1043                 || rb_tree_count_black(rbt->rbt_root));
1044
1045         /*
1046          * The root must be black.
1047          * There can never be two adjacent red nodes.
1048          */
1049         if (red_check) {
1050                 KASSERT(rbt->rbt_root == NULL || RB_BLACK_P(rbt->rbt_root));
1051                 TAILQ_FOREACH(self, &rbt->rbt_nodes, rb_link) {
1052                         rb_tree_check_node(rbt, self, NULL, true);
1053                 }
1054         }
1055 }
1056 #endif /* RBDEBUG */