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