util/redblack.c

   1 /* Copyright (c) 2007-2014 Massachusetts Institute of Technology
   2  *
   3  * Permission is hereby granted, free of charge, to any person obtaining
   4  * a copy of this software and associated documentation files (the
   5  * "Software"), to deal in the Software without restriction, including
   6  * without limitation the rights to use, copy, modify, merge, publish,
   7  * distribute, sublicense, and/or sell copies of the Software, and to
   8  * permit persons to whom the Software is furnished to do so, subject to
   9  * the following conditions:
  10  *
  11  * The above copyright notice and this permission notice shall be
  12  * included in all copies or substantial portions of the Software.
  13  *
  14  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  15  * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  16  * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  17  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
  18  * LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
  19  * OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
  20  * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  21  */
  22
  23 /* simple implementation of red-black trees optimized for use with DIRECT */
  24
  25 #include <stddef.h>
  26 #include <stdlib.h>
  27 #include "redblack.h"
  28
  29 /* it is convenient to use an explicit node for NULL nodes ... we need
  30    to be careful never to change this node indirectly via one of our
  31    pointers!  */
  32 rb_node nil = {&nil, &nil, &nil, 0, BLACK};
  33 #define NIL (&nil)
  34
  35 void rb_tree_init(rb_tree *t, rb_compare compare) {
  36      t->compare = compare;
  37      t->root = NIL;
  38      t->N = 0;
  39 }
  40
  41 static void destroy(rb_node *n)
  42 {
  43      if (n != NIL) {
  44           destroy(n->l); destroy(n->r);
  45           free(n);
  46      }
  47 }
  48
  49 void rb_tree_destroy(rb_tree *t)
  50 {
  51      destroy(t->root);
  52      t->root = NIL;
  53 }
  54
  55 void rb_tree_destroy_with_keys(rb_tree *t)
  56 {
  57      rb_node *n = rb_tree_min(t);
  58      while (n) {
  59           free(n->k); n->k = NULL;
  60           n = rb_tree_succ(n);
  61      }
  62      rb_tree_destroy(t);
  63 }
  64
  65 static void rotate_left(rb_node *p, rb_tree *t)
  66 {
  67      rb_node *n = p->r; /* must be non-NIL */
  68      p->r = n->l;
  69      n->l = p;
  70      if (p->p != NIL) {
  71           if (p == p->p->l) p->p->l = n;
  72           else p->p->r = n;
  73      }
  74      else
  75           t->root = n;
  76      n->p = p->p;
  77      p->p = n;
  78      if (p->r != NIL) p->r->p = p;
  79 }
  80
  81 static void rotate_right(rb_node *p, rb_tree *t)
  82 {
  83      rb_node *n = p->l; /* must be non-NIL */
  84      p->l = n->r;
  85      n->r = p;
  86      if (p->p != NIL) {
  87           if (p == p->p->l) p->p->l = n;
  88           else p->p->r = n;
  89      }
  90      else
  91           t->root = n;
  92      n->p = p->p;
  93      p->p = n;
  94      if (p->l != NIL) p->l->p = p;
  95 }
  96
  97 static void insert_node(rb_tree *t, rb_node *n)
  98 {
  99      rb_compare compare = t->compare;
 100      rb_key k = n->k;
 101      rb_node *p = t->root;
 102      n->c = RED;
 103      n->p = n->l = n->r = NIL;
 104      t->N++;
 105      if (p == NIL) {
 106           t->root = n;
 107           n->c = BLACK;
 108           return;
 109      }
 110      /* insert (RED) node into tree */
 111      while (1) {
 112           if (compare(k, p->k) <= 0) { /* k <= p->k */
 113                if (p->l != NIL)
 114                     p = p->l;
 115                else {
 116                     p->l = n;
 117                     n->p = p;
 118                     break;
 119                }
 120           }
 121           else {
 122                if (p->r != NIL)
 123                     p = p->r;
 124                else {
 125                     p->r = n;
 126                     n->p = p;
 127                     break;
 128                }
 129           }
 130      }
 131  fixtree:
 132      if (n->p->c == RED) { /* red cannot have red child */
 133           rb_node *u = p == p->p->l ? p->p->r : p->p->l;
 134           if (u != NIL && u->c == RED) {
 135                p->c = u->c = BLACK;
 136                n = p->p;
 137                if ((p = n->p) != NIL) {
 138                     n->c = RED;
 139                     goto fixtree;
 140                }
 141           }
 142           else {
 143                if (n == p->r && p == p->p->l) {
 144                     rotate_left(p, t);
 145                     p = n; n = n->l;
 146                }
 147                else if (n == p->l && p == p->p->r) {
 148                     rotate_right(p, t);
 149                     p = n; n = n->r;
 150                }
 151                p->c = BLACK;
 152                p->p->c = RED;
 153                if (n == p->l && p == p->p->l)
 154                     rotate_right(p->p, t);
 155                else if (n == p->r && p == p->p->r)
 156                     rotate_left(p->p, t);
 157           }
 158
 159      }
 160 }
 161
 162 rb_node *rb_tree_insert(rb_tree *t, rb_key k)
 163 {
 164      rb_node *n = (rb_node *) malloc(sizeof(rb_node));
 165      if (!n) return NULL;
 166      n->k = k;
 167      insert_node(t, n);
 168      return n;
 169 }
 170
 171 static int check_node(rb_node *n, int *nblack, rb_tree *t)
 172 {
 173      int nbl, nbr;
 174      rb_compare compare = t->compare;
 175      if (n == NIL) { *nblack = 0; return 1; }
 176      if (n->r != NIL && n->r->p != n) return 0;
 177      if (n->r != NIL && compare(n->r->k, n->k) < 0)
 178           return 0;
 179      if (n->l != NIL && n->l->p != n) return 0;
 180      if (n->l != NIL && compare(n->l->k, n->k) > 0)
 181           return 0;
 182      if (n->c == RED) {
 183           if (n->r != NIL && n->r->c == RED) return 0;
 184           if (n->l != NIL && n->l->c == RED) return 0;
 185      }
 186      if (!(check_node(n->r, &nbl, t) && check_node(n->l, &nbr, t)))
 187           return 0;
 188      if (nbl != nbr) return 0;
 189      *nblack = nbl + (n->c == BLACK);
 190      return 1;
 191 }
 192 int rb_tree_check(rb_tree *t)
 193 {
 194      int nblack;
 195      if (nil.c != BLACK) return 0;
 196      if (nil.p != NIL || nil.r != NIL || nil.l != NIL) return 0;
 197      if (t->root == NIL) return 1;
 198      if (t->root->c != BLACK) return 0;
 199      return check_node(t->root, &nblack, t);
 200 }
 201
 202 rb_node *rb_tree_find(rb_tree *t, rb_key k)
 203 {
 204      rb_compare compare = t->compare;
 205      rb_node *p = t->root;
 206      while (p != NIL) {
 207           int comp = compare(k, p->k);
 208           if (!comp) return p;
 209           p = comp <= 0 ? p->l : p->r;
 210      }
 211      return NULL;
 212 }
 213
 214 /* find greatest point in subtree p that is <= k */
 215 static rb_node *find_le(rb_node *p, rb_key k, rb_tree *t)
 216 {
 217      rb_compare compare = t->compare;
 218      while (p != NIL) {
 219           if (compare(p->k, k) <= 0) { /* p->k <= k */
 220                rb_node *r = find_le(p->r, k, t);
 221                if (r) return r;
 222                else return p;
 223           }
 224           else /* p->k > k */
 225                p = p->l;
 226      }
 227      return NULL; /* k < everything in subtree */
 228 }
 229
 230 /* find greatest point in t <= k */
 231 rb_node *rb_tree_find_le(rb_tree *t, rb_key k)
 232 {
 233      return find_le(t->root, k, t);
 234 }
 235
 236 /* find greatest point in subtree p that is < k */
 237 static rb_node *find_lt(rb_node *p, rb_key k, rb_tree *t)
 238 {
 239      rb_compare compare = t->compare;
 240      while (p != NIL) {
 241           if (compare(p->k, k) < 0) { /* p->k < k */
 242                rb_node *r = find_lt(p->r, k, t);
 243                if (r) return r;
 244                else return p;
 245           }
 246           else /* p->k >= k */
 247                p = p->l;
 248      }
 249      return NULL; /* k <= everything in subtree */
 250 }
 251
 252 /* find greatest point in t < k */
 253 rb_node *rb_tree_find_lt(rb_tree *t, rb_key k)
 254 {
 255      return find_lt(t->root, k, t);
 256 }
 257
 258 /* find least point in subtree p that is > k */
 259 static rb_node *find_gt(rb_node *p, rb_key k, rb_tree *t)
 260 {
 261      rb_compare compare = t->compare;
 262      while (p != NIL) {
 263           if (compare(p->k, k) > 0) { /* p->k > k */
 264                rb_node *l = find_gt(p->l, k, t);
 265                if (l) return l;
 266                else return p;
 267           }
 268           else /* p->k <= k */
 269                p = p->r;
 270      }
 271      return NULL; /* k >= everything in subtree */
 272 }
 273
 274 /* find least point in t > k */
 275 rb_node *rb_tree_find_gt(rb_tree *t, rb_key k)
 276 {
 277      return find_gt(t->root, k, t);
 278 }
 279
 280 rb_node *rb_tree_min(rb_tree *t)
 281 {
 282      rb_node *n = t->root;
 283      while (n != NIL && n->l != NIL)
 284           n = n->l;
 285      return(n == NIL ? NULL : n);
 286 }
 287
 288 rb_node *rb_tree_max(rb_tree *t)
 289 {
 290      rb_node *n = t->root;
 291      while (n != NIL && n->r != NIL)
 292           n = n->r;
 293      return(n == NIL ? NULL : n);
 294 }
 295
 296 rb_node *rb_tree_succ(rb_node *n)
 297 {
 298      if (!n) return NULL;
 299      if (n->r == NIL) {
 300           rb_node *prev;
 301           do {
 302                prev = n;
 303                n = n->p;
 304           } while (prev == n->r && n != NIL);
 305           return n == NIL ? NULL : n;
 306      }
 307      else {
 308           n = n->r;
 309           while (n->l != NIL)
 310                n = n->l;
 311           return n;
 312      }
 313 }
 314
 315 rb_node *rb_tree_pred(rb_node *n)
 316 {
 317      if (!n) return NULL;
 318      if (n->l == NIL) {
 319           rb_node *prev;
 320           do {
 321                prev = n;
 322                n = n->p;
 323           } while (prev == n->l && n != NIL);
 324           return n == NIL ? NULL : n;
 325      }
 326      else {
 327           n = n->l;
 328           while (n->r != NIL)
 329                n = n->r;
 330           return n;
 331      }
 332 }
 333
 334 rb_node *rb_tree_remove(rb_tree *t, rb_node *n)
 335 {
 336      rb_key k = n->k;
 337      rb_node *m, *mp;
 338      if (n->l != NIL && n->r != NIL) {
 339           rb_node *lmax = n->l;
 340           while (lmax->r != NIL) lmax = lmax->r;
 341           n->k = lmax->k;
 342           n = lmax;
 343      }
 344      m = n->l != NIL ? n->l : n->r;
 345      if (n->p != NIL) {
 346           if (n->p->r == n) n->p->r = m;
 347           else n->p->l = m;
 348      }
 349      else
 350           t->root = m;
 351      mp = n->p;
 352      if (m != NIL) m->p = mp;
 353      if (n->c == BLACK) {
 354           if (m->c == RED)
 355                m->c = BLACK;
 356           else {
 357           deleteblack:
 358                if (mp != NIL) {
 359                     rb_node *s = m == mp->l ? mp->r : mp->l;
 360                     if (s->c == RED) {
 361                          mp->c = RED;
 362                          s->c = BLACK;
 363                          if (m == mp->l) rotate_left(mp, t);
 364                          else rotate_right(mp, t);
 365                          s = m == mp->l ? mp->r : mp->l;
 366                     }
 367                     if (mp->c == BLACK && s->c == BLACK
 368                         && s->l->c == BLACK && s->r->c == BLACK) {
 369                          if (s != NIL) s->c = RED;
 370                          m = mp; mp = m->p;
 371                          goto deleteblack;
 372                     }
 373                     else if (mp->c == RED && s->c == BLACK &&
 374                              s->l->c == BLACK && s->r->c == BLACK) {
 375                          if (s != NIL) s->c = RED;
 376                          mp->c = BLACK;
 377                     }
 378                     else {
 379                          if (m == mp->l && s->c == BLACK &&
 380                              s->l->c == RED && s->r->c == BLACK) {
 381                               s->c = RED;
 382                               s->l->c = BLACK;
 383                               rotate_right(s, t);
 384                               s = m == mp->l ? mp->r : mp->l;
 385                          }
 386                          else if (m == mp->r && s->c == BLACK &&
 387                                   s->r->c == RED && s->l->c == BLACK) {
 388                               s->c = RED;
 389                               s->r->c = BLACK;
 390                               rotate_left(s, t);
 391                               s = m == mp->l ? mp->r : mp->l;
 392                          }
 393                          s->c = mp->c;
 394                          mp->c = BLACK;
 395                          if (m == mp->l) {
 396                               s->r->c = BLACK;
 397                               rotate_left(mp, t);
 398                          }
 399                          else {
 400                               s->l->c = BLACK;
 401                               rotate_right(mp, t);
 402                          }
 403                     }
 404                }
 405           }
 406      }
 407      t->N--;
 408      n->k = k; /* n may have changed during remove */
 409      return n; /* the node that was deleted may be different from initial n */
 410 }
 411
 412 rb_node *rb_tree_resort(rb_tree *t, rb_node *n)
 413 {
 414      n = rb_tree_remove(t, n);
 415      insert_node(t, n);
 416      return n;
 417 }
 418
 419 /* shift all key pointers by kshift ... this is useful when the keys
 420    are pointers into another array, that has been resized with realloc */
 421 static void shift_keys(rb_node *n, ptrdiff_t kshift) /* assumes n != NIL */
 422 {
 423      n->k += kshift;
 424      if (n->l != NIL) shift_keys(n->l, kshift);
 425      if (n->r != NIL) shift_keys(n->r, kshift);
 426 }
 427 void rb_tree_shift_keys(rb_tree *t, ptrdiff_t kshift)
 428 {
 429      if (t->root != NIL) shift_keys(t->root, kshift);
 430 }