src/util/rb_tree.c

   1 /*
   2  * Copyright © 2017 Jason Ekstrand
   3  *
   4  * Permission is hereby granted, free of charge, to any person obtaining a
   5  * copy of this software and associated documentation files (the "Software"),
   6  * to deal in the Software without restriction, including without limitation
   7  * the rights to use, copy, modify, merge, publish, distribute, sublicense,
   8  * and/or sell copies of the Software, and to permit persons to whom the
   9  * Software is furnished to do so, subject to the following conditions:
  10  *
  11  * The above copyright notice and this permission notice shall be included in
  12  * all copies or substantial portions of the Software.
  13  *
  14  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  15  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  16  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  17  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  18  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  19  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  20  * DEALINGS IN THE SOFTWARE.
  21  */
  22
  23 #include "rb_tree.h"
  24
  25 /** \file rb_tree.c
  26  *
  27  * An implementation of a red-black tree
  28  *
  29  * This file implements the guts of a red-black tree.  The implementation
  30  * is mostly based on the one in "Introduction to Algorithms", third
  31  * edition, by Cormen, Leiserson, Rivest, and Stein.  The primary
  32  * divergence in our algorithms from those presented in CLRS is that we use
  33  * NULL for the leaves instead of a sentinel.  This means we have to do a
  34  * tiny bit more tracking in our implementation of delete but it makes the
  35  * algorithms far more explicit than stashing stuff in the sentinel.
  36  */
  37
  38 #include <stdlib.h>
  39 #include <string.h>
  40 #include <assert.h>
  41
  42 static bool
  43 rb_node_is_black(struct rb_node *n)
  44 {
  45     /* NULL nodes are leaves and therefore black */
  46     return (n == NULL) || (n->parent & 1);
  47 }
  48
  49 static bool
  50 rb_node_is_red(struct rb_node *n)
  51 {
  52     return !rb_node_is_black(n);
  53 }
  54
  55 static void
  56 rb_node_set_black(struct rb_node *n)
  57 {
  58     n->parent |= 1;
  59 }
  60
  61 static void
  62 rb_node_set_red(struct rb_node *n)
  63 {
  64     n->parent &= ~1ull;
  65 }
  66
  67 static void
  68 rb_node_copy_color(struct rb_node *dst, struct rb_node *src)
  69 {
  70     dst->parent = (dst->parent & ~1ull) | (src->parent & 1);
  71 }
  72
  73 static void
  74 rb_node_set_parent(struct rb_node *n, struct rb_node *p)
  75 {
  76     n->parent = (n->parent & 1) | (uintptr_t)p;
  77 }
  78
  79 static struct rb_node *
  80 rb_node_minimum(struct rb_node *node)
  81 {
  82     while (node->left)
  83         node = node->left;
  84     return node;
  85 }
  86
  87 static struct rb_node *
  88 rb_node_maximum(struct rb_node *node)
  89 {
  90     while (node->right)
  91         node = node->right;
  92     return node;
  93 }
  94
  95 void
  96 rb_tree_init(struct rb_tree *T)
  97 {
  98     T->root = NULL;
  99 }
 100
 101 /**
 102  * Replace the subtree of T rooted at u with the subtree rooted at v
 103  *
 104  * This is called RB-transplant in CLRS.
 105  *
 106  * The node to be replaced is assumed to be a non-leaf.
 107  */
 108 static void
 109 rb_tree_splice(struct rb_tree *T, struct rb_node *u, struct rb_node *v)
 110 {
 111     assert(u);
 112     struct rb_node *p = rb_node_parent(u);
 113     if (p == NULL) {
 114         assert(T->root == u);
 115         T->root = v;
 116     } else if (u == p->left) {
 117         p->left = v;
 118     } else {
 119         assert(u == p->right);
 120         p->right = v;
 121     }
 122     if (v)
 123         rb_node_set_parent(v, p);
 124 }
 125
 126 static void
 127 rb_tree_rotate_left(struct rb_tree *T, struct rb_node *x)
 128 {
 129     assert(x && x->right);
 130
 131     struct rb_node *y = x->right;
 132     x->right = y->left;
 133     if (y->left)
 134         rb_node_set_parent(y->left, x);
 135     rb_tree_splice(T, x, y);
 136     y->left = x;
 137     rb_node_set_parent(x, y);
 138 }
 139
 140 static void
 141 rb_tree_rotate_right(struct rb_tree *T, struct rb_node *y)
 142 {
 143     assert(y && y->left);
 144
 145     struct rb_node *x = y->left;
 146     y->left = x->right;
 147     if (x->right)
 148         rb_node_set_parent(x->right, y);
 149     rb_tree_splice(T, y, x);
 150     x->right = y;
 151     rb_node_set_parent(y, x);
 152 }
 153
 154 void
 155 rb_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
 156                   struct rb_node *node, bool insert_left)
 157 {
 158     /* This sets null children, parent, and a color of red */
 159     memset(node, 0, sizeof(*node));
 160
 161     if (parent == NULL) {
 162         assert(T->root == NULL);
 163         T->root = node;
 164         rb_node_set_black(node);
 165         return;
 166     }
 167
 168     if (insert_left) {
 169         assert(parent->left == NULL);
 170         parent->left = node;
 171     } else {
 172         assert(parent->right == NULL);
 173         parent->right = node;
 174     }
 175     rb_node_set_parent(node, parent);
 176
 177     /* Now we do the insertion fixup */
 178     struct rb_node *z = node;
 179     while (rb_node_is_red(rb_node_parent(z))) {
 180         struct rb_node *z_p = rb_node_parent(z);
 181         assert(z == z_p->left || z == z_p->right);
 182         struct rb_node *z_p_p = rb_node_parent(z_p);
 183         assert(z_p_p != NULL);
 184         if (z_p == z_p_p->left) {
 185             struct rb_node *y = z_p_p->right;
 186             if (rb_node_is_red(y)) {
 187                 rb_node_set_black(z_p);
 188                 rb_node_set_black(y);
 189                 rb_node_set_red(z_p_p);
 190                 z = z_p_p;
 191             } else {
 192                 if (z == z_p->right) {
 193                     z = z_p;
 194                     rb_tree_rotate_left(T, z);
 195                     /* We changed z */
 196                     z_p = rb_node_parent(z);
 197                     assert(z == z_p->left || z == z_p->right);
 198                     z_p_p = rb_node_parent(z_p);
 199                 }
 200                 rb_node_set_black(z_p);
 201                 rb_node_set_red(z_p_p);
 202                 rb_tree_rotate_right(T, z_p_p);
 203             }
 204         } else {
 205             struct rb_node *y = z_p_p->left;
 206             if (rb_node_is_red(y)) {
 207                 rb_node_set_black(z_p);
 208                 rb_node_set_black(y);
 209                 rb_node_set_red(z_p_p);
 210                 z = z_p_p;
 211             } else {
 212                 if (z == z_p->left) {
 213                     z = z_p;
 214                     rb_tree_rotate_right(T, z);
 215                     /* We changed z */
 216                     z_p = rb_node_parent(z);
 217                     assert(z == z_p->left || z == z_p->right);
 218                     z_p_p = rb_node_parent(z_p);
 219                 }
 220                 rb_node_set_black(z_p);
 221                 rb_node_set_red(z_p_p);
 222                 rb_tree_rotate_left(T, z_p_p);
 223             }
 224         }
 225     }
 226     rb_node_set_black(T->root);
 227 }
 228
 229 void
 230 rb_tree_remove(struct rb_tree *T, struct rb_node *z)
 231 {
 232     /* x_p is always the parent node of X.  We have to track this
 233      * separately because x may be NULL.
 234      */
 235     struct rb_node *x, *x_p;
 236     struct rb_node *y = z;
 237     bool y_was_black = rb_node_is_black(y);
 238     if (z->left == NULL) {
 239         x = z->right;
 240         x_p = rb_node_parent(z);
 241         rb_tree_splice(T, z, x);
 242     } else if (z->right == NULL) {
 243         x = z->left;
 244         x_p = rb_node_parent(z);
 245         rb_tree_splice(T, z, x);
 246     } else {
 247         /* Find the minimum sub-node of z->right */
 248         y = rb_node_minimum(z->right);
 249         y_was_black = rb_node_is_black(y);
 250
 251         x = y->right;
 252         if (rb_node_parent(y) == z) {
 253             x_p = y;
 254         } else {
 255             x_p = rb_node_parent(y);
 256             rb_tree_splice(T, y, x);
 257             y->right = z->right;
 258             rb_node_set_parent(y->right, y);
 259         }
 260         assert(y->left == NULL);
 261         rb_tree_splice(T, z, y);
 262         y->left = z->left;
 263         rb_node_set_parent(y->left, y);
 264         rb_node_copy_color(y, z);
 265     }
 266
 267     assert(x_p == NULL || x == x_p->left || x == x_p->right);
 268
 269     if (!y_was_black)
 270         return;
 271
 272     /* Fixup RB tree after the delete */
 273     while (x != T->root && rb_node_is_black(x)) {
 274         if (x == x_p->left) {
 275             struct rb_node *w = x_p->right;
 276             if (rb_node_is_red(w)) {
 277                 rb_node_set_black(w);
 278                 rb_node_set_red(x_p);
 279                 rb_tree_rotate_left(T, x_p);
 280                 assert(x == x_p->left);
 281                 w = x_p->right;
 282             }
 283             if (rb_node_is_black(w->left) && rb_node_is_black(w->right)) {
 284                 rb_node_set_red(w);
 285                 x = x_p;
 286             } else {
 287                 if (rb_node_is_black(w->right)) {
 288                     rb_node_set_black(w->left);
 289                     rb_node_set_red(w);
 290                     rb_tree_rotate_right(T, w);
 291                     w = x_p->right;
 292                 }
 293                 rb_node_copy_color(w, x_p);
 294                 rb_node_set_black(x_p);
 295                 rb_node_set_black(w->right);
 296                 rb_tree_rotate_left(T, x_p);
 297                 x = T->root;
 298             }
 299         } else {
 300             struct rb_node *w = x_p->left;
 301             if (rb_node_is_red(w)) {
 302                 rb_node_set_black(w);
 303                 rb_node_set_red(x_p);
 304                 rb_tree_rotate_right(T, x_p);
 305                 assert(x == x_p->right);
 306                 w = x_p->left;
 307             }
 308             if (rb_node_is_black(w->right) && rb_node_is_black(w->left)) {
 309                 rb_node_set_red(w);
 310                 x = x_p;
 311             } else {
 312                 if (rb_node_is_black(w->left)) {
 313                     rb_node_set_black(w->right);
 314                     rb_node_set_red(w);
 315                     rb_tree_rotate_left(T, w);
 316                     w = x_p->left;
 317                 }
 318                 rb_node_copy_color(w, x_p);
 319                 rb_node_set_black(x_p);
 320                 rb_node_set_black(w->left);
 321                 rb_tree_rotate_right(T, x_p);
 322                 x = T->root;
 323             }
 324         }
 325         x_p = rb_node_parent(x);
 326     }
 327     if (x)
 328         rb_node_set_black(x);
 329 }
 330
 331 struct rb_node *
 332 rb_tree_first(struct rb_tree *T)
 333 {
 334     return T->root ? rb_node_minimum(T->root) : NULL;
 335 }
 336
 337 struct rb_node *
 338 rb_tree_last(struct rb_tree *T)
 339 {
 340     return T->root ? rb_node_maximum(T->root) : NULL;
 341 }
 342
 343 struct rb_node *
 344 rb_node_next(struct rb_node *node)
 345 {
 346     if (node->right) {
 347         /* If we have a right child, then the next thing (compared to this
 348          * node) is the left-most child of our right child.
 349          */
 350         return rb_node_minimum(node->right);
 351     } else {
 352         /* If node doesn't have a right child, crawl back up the to the
 353          * left until we hit a parent to the right.
 354          */
 355         struct rb_node *p = rb_node_parent(node);
 356         while (p && node == p->right) {
 357             node = p;
 358             p = rb_node_parent(node);
 359         }
 360         assert(p == NULL || node == p->left);
 361         return p;
 362     }
 363 }
 364
 365 struct rb_node *
 366 rb_node_prev(struct rb_node *node)
 367 {
 368     if (node->left) {
 369         /* If we have a left child, then the previous thing (compared to
 370          * this node) is the right-most child of our left child.
 371          */
 372         return rb_node_maximum(node->left);
 373     } else {
 374         /* If node doesn't have a left child, crawl back up the to the
 375          * right until we hit a parent to the left.
 376          */
 377         struct rb_node *p = rb_node_parent(node);
 378         while (p && node == p->left) {
 379             node = p;
 380             p = rb_node_parent(node);
 381         }
 382         assert(p == NULL || node == p->right);
 383         return p;
 384     }
 385 }
 386
 387 static void
 388 validate_rb_node(struct rb_node *n, int black_depth)
 389 {
 390     if (n == NULL) {
 391         assert(black_depth == 0);
 392         return;
 393     }
 394
 395     if (rb_node_is_black(n)) {
 396         black_depth--;
 397     } else {
 398         assert(rb_node_is_black(n->left));
 399         assert(rb_node_is_black(n->right));
 400     }
 401
 402     validate_rb_node(n->left, black_depth);
 403     validate_rb_node(n->right, black_depth);
 404 }
 405
 406 void
 407 rb_tree_validate(struct rb_tree *T)
 408 {
 409     if (T->root == NULL)
 410         return;
 411
 412     assert(rb_node_is_black(T->root));
 413
 414     unsigned black_depth = 0;
 415     for (struct rb_node *n = T->root; n; n = n->left) {
 416         if (rb_node_is_black(n))
 417             black_depth++;
 418     }
 419
 420     validate_rb_node(T->root, black_depth);
 421 }