src/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
  34 #define NIL (&nil)
  35
  36 void rb_tree_init(rb_tree * t, rb_compare compare)
  37 {
  38     t->compare = compare;
  39     t->root = NIL;
  40     t->N = 0;
  41 }
  42
  43 static void destroy(rb_node * n)
  44 {
  45     if (n != NIL) {
  46         destroy(n->l);
  47         destroy(n->r);
  48         free(n);
  49     }
  50 }
  51
  52 void rb_tree_destroy(rb_tree * t)
  53 {
  54     destroy(t->root);
  55     t->root = NIL;
  56 }
  57
  58 void rb_tree_destroy_with_keys(rb_tree * t)
  59 {
  60     rb_node *n = rb_tree_min(t);
  61     while (n) {
  62         free(n->k);
  63         n->k = NULL;
  64         n = rb_tree_succ(n);
  65     }
  66     rb_tree_destroy(t);
  67 }
  68
  69 static void rotate_left(rb_node * p, rb_tree * t)
  70 {
  71     rb_node *n = p->r;          /* must be non-NIL */
  72     p->r = n->l;
  73     n->l = p;
  74     if (p->p != NIL) {
  75         if (p == p->p->l)
  76             p->p->l = n;
  77         else
  78             p->p->r = n;
  79     } else
  80         t->root = n;
  81     n->p = p->p;
  82     p->p = n;
  83     if (p->r != NIL)
  84         p->r->p = p;
  85 }
  86
  87 static void rotate_right(rb_node * p, rb_tree * t)
  88 {
  89     rb_node *n = p->l;          /* must be non-NIL */
  90     p->l = n->r;
  91     n->r = p;
  92     if (p->p != NIL) {
  93         if (p == p->p->l)
  94             p->p->l = n;
  95         else
  96             p->p->r = n;
  97     } else
  98         t->root = n;
  99     n->p = p->p;
 100     p->p = n;
 101     if (p->l != NIL)
 102         p->l->p = p;
 103 }
 104
 105 static void insert_node(rb_tree * t, rb_node * n)
 106 {
 107     rb_compare compare = t->compare;
 108     rb_key k = n->k;
 109     rb_node *p = t->root;
 110     n->c = RED;
 111     n->p = n->l = n->r = NIL;
 112     t->N++;
 113     if (p == NIL) {
 114         t->root = n;
 115         n->c = BLACK;
 116         return;
 117     }
 118     /* insert (RED) node into tree */
 119     while (1) {
 120         if (compare(k, p->k) <= 0) {    /* k <= p->k */
 121             if (p->l != NIL)
 122                 p = p->l;
 123             else {
 124                 p->l = n;
 125                 n->p = p;
 126                 break;
 127             }
 128         } else {
 129             if (p->r != NIL)
 130                 p = p->r;
 131             else {
 132                 p->r = n;
 133                 n->p = p;
 134                 break;
 135             }
 136         }
 137     }
 138   fixtree:
 139     if (n->p->c == RED) {       /* red cannot have red child */
 140         rb_node *u = p == p->p->l ? p->p->r : p->p->l;
 141         if (u != NIL && u->c == RED) {
 142             p->c = u->c = BLACK;
 143             n = p->p;
 144             if ((p = n->p) != NIL) {
 145                 n->c = RED;
 146                 goto fixtree;
 147             }
 148         } else {
 149             if (n == p->r && p == p->p->l) {
 150                 rotate_left(p, t);
 151                 p = n;
 152                 n = n->l;
 153             } else if (n == p->l && p == p->p->r) {
 154                 rotate_right(p, t);
 155                 p = n;
 156                 n = n->r;
 157             }
 158             p->c = BLACK;
 159             p->p->c = RED;
 160             if (n == p->l && p == p->p->l)
 161                 rotate_right(p->p, t);
 162             else if (n == p->r && p == p->p->r)
 163                 rotate_left(p->p, t);
 164         }
 165
 166     }
 167 }
 168
 169 rb_node *rb_tree_insert(rb_tree * t, rb_key k)
 170 {
 171     rb_node *n = (rb_node *) malloc(sizeof(rb_node));
 172     if (!n)
 173         return NULL;
 174     n->k = k;
 175     insert_node(t, n);
 176     return n;
 177 }
 178
 179 static int check_node(rb_node * n, int *nblack, rb_tree * t)
 180 {
 181     int nbl, nbr;
 182     rb_compare compare = t->compare;
 183     if (n == NIL) {
 184         *nblack = 0;
 185         return 1;
 186     }
 187     if (n->r != NIL && n->r->p != n)
 188         return 0;
 189     if (n->r != NIL && compare(n->r->k, n->k) < 0)
 190         return 0;
 191     if (n->l != NIL && n->l->p != n)
 192         return 0;
 193     if (n->l != NIL && compare(n->l->k, n->k) > 0)
 194         return 0;
 195     if (n->c == RED) {
 196         if (n->r != NIL && n->r->c == RED)
 197             return 0;
 198         if (n->l != NIL && n->l->c == RED)
 199             return 0;
 200     }
 201     if (!(check_node(n->r, &nbl, t) && check_node(n->l, &nbr, t)))
 202         return 0;
 203     if (nbl != nbr)
 204         return 0;
 205     *nblack = nbl + (n->c == BLACK);
 206     return 1;
 207 }
 208
 209 int rb_tree_check(rb_tree * t)
 210 {
 211     int nblack;
 212     if (nil.c != BLACK)
 213         return 0;
 214     if (nil.p != NIL || nil.r != NIL || nil.l != NIL)
 215         return 0;
 216     if (t->root == NIL)
 217         return 1;
 218     if (t->root->c != BLACK)
 219         return 0;
 220     return check_node(t->root, &nblack, t);
 221 }
 222
 223 rb_node *rb_tree_find(rb_tree * t, rb_key k)
 224 {
 225     rb_compare compare = t->compare;
 226     rb_node *p = t->root;
 227     while (p != NIL) {
 228         int comp = compare(k, p->k);
 229         if (!comp)
 230             return p;
 231         p = comp <= 0 ? p->l : p->r;
 232     }
 233     return NULL;
 234 }
 235
 236 /* find greatest point in subtree p that is <= k */
 237 static rb_node *find_le(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_le(p->r, k, t);
 243             if (r)
 244                 return r;
 245             else
 246                 return p;
 247         } else                  /* p->k > k */
 248             p = p->l;
 249     }
 250     return NULL;                /* k < everything in subtree */
 251 }
 252
 253 /* find greatest point in t <= k */
 254 rb_node *rb_tree_find_le(rb_tree * t, rb_key k)
 255 {
 256     return find_le(t->root, k, t);
 257 }
 258
 259 /* find greatest point in subtree p that is < k */
 260 static rb_node *find_lt(rb_node * p, rb_key k, rb_tree * t)
 261 {
 262     rb_compare compare = t->compare;
 263     while (p != NIL) {
 264         if (compare(p->k, k) < 0) {     /* p->k < k */
 265             rb_node *r = find_lt(p->r, k, t);
 266             if (r)
 267                 return r;
 268             else
 269                 return p;
 270         } else                  /* p->k >= k */
 271             p = p->l;
 272     }
 273     return NULL;                /* k <= everything in subtree */
 274 }
 275
 276 /* find greatest point in t < k */
 277 rb_node *rb_tree_find_lt(rb_tree * t, rb_key k)
 278 {
 279     return find_lt(t->root, k, t);
 280 }
 281
 282 /* find least point in subtree p that is > k */
 283 static rb_node *find_gt(rb_node * p, rb_key k, rb_tree * t)
 284 {
 285     rb_compare compare = t->compare;
 286     while (p != NIL) {
 287         if (compare(p->k, k) > 0) {     /* p->k > k */
 288             rb_node *l = find_gt(p->l, k, t);
 289             if (l)
 290                 return l;
 291             else
 292                 return p;
 293         } else                  /* p->k <= k */
 294             p = p->r;
 295     }
 296     return NULL;                /* k >= everything in subtree */
 297 }
 298
 299 /* find least point in t > k */
 300 rb_node *rb_tree_find_gt(rb_tree * t, rb_key k)
 301 {
 302     return find_gt(t->root, k, t);
 303 }
 304
 305 rb_node *rb_tree_min(rb_tree * t)
 306 {
 307     rb_node *n = t->root;
 308     while (n != NIL && n->l != NIL)
 309         n = n->l;
 310     return (n == NIL ? NULL : n);
 311 }
 312
 313 rb_node *rb_tree_max(rb_tree * t)
 314 {
 315     rb_node *n = t->root;
 316     while (n != NIL && n->r != NIL)
 317         n = n->r;
 318     return (n == NIL ? NULL : n);
 319 }
 320
 321 rb_node *rb_tree_succ(rb_node * n)
 322 {
 323     if (!n)
 324         return NULL;
 325     if (n->r == NIL) {
 326         rb_node *prev;
 327         do {
 328             prev = n;
 329             n = n->p;
 330         } while (prev == n->r && n != NIL);
 331         return n == NIL ? NULL : n;
 332     } else {
 333         n = n->r;
 334         while (n->l != NIL)
 335             n = n->l;
 336         return n;
 337     }
 338 }
 339
 340 rb_node *rb_tree_pred(rb_node * n)
 341 {
 342     if (!n)
 343         return NULL;
 344     if (n->l == NIL) {
 345         rb_node *prev;
 346         do {
 347             prev = n;
 348             n = n->p;
 349         } while (prev == n->l && n != NIL);
 350         return n == NIL ? NULL : n;
 351     } else {
 352         n = n->l;
 353         while (n->r != NIL)
 354             n = n->r;
 355         return n;
 356     }
 357 }
 358
 359 rb_node *rb_tree_remove(rb_tree * t, rb_node * n)
 360 {
 361     rb_key k = n->k;
 362     rb_node *m, *mp;
 363     if (n->l != NIL && n->r != NIL) {
 364         rb_node *lmax = n->l;
 365         while (lmax->r != NIL)
 366             lmax = lmax->r;
 367         n->k = lmax->k;
 368         n = lmax;
 369     }
 370     m = n->l != NIL ? n->l : n->r;
 371     if (n->p != NIL) {
 372         if (n->p->r == n)
 373             n->p->r = m;
 374         else
 375             n->p->l = m;
 376     } else
 377         t->root = m;
 378     mp = n->p;
 379     if (m != NIL)
 380         m->p = mp;
 381     if (n->c == BLACK) {
 382         if (m->c == RED)
 383             m->c = BLACK;
 384         else {
 385           deleteblack:
 386             if (mp != NIL) {
 387                 rb_node *s = m == mp->l ? mp->r : mp->l;
 388                 if (s->c == RED) {
 389                     mp->c = RED;
 390                     s->c = BLACK;
 391                     if (m == mp->l)
 392                         rotate_left(mp, t);
 393                     else
 394                         rotate_right(mp, t);
 395                     s = m == mp->l ? mp->r : mp->l;
 396                 }
 397                 if (mp->c == BLACK && s->c == BLACK && s->l->c == BLACK && s->r->c == BLACK) {
 398                     if (s != NIL)
 399                         s->c = RED;
 400                     m = mp;
 401                     mp = m->p;
 402                     goto deleteblack;
 403                 } else if (mp->c == RED && s->c == BLACK && s->l->c == BLACK && s->r->c == BLACK) {
 404                     if (s != NIL)
 405                         s->c = RED;
 406                     mp->c = BLACK;
 407                 } else {
 408                     if (m == mp->l && s->c == BLACK && s->l->c == RED && s->r->c == BLACK) {
 409                         s->c = RED;
 410                         s->l->c = BLACK;
 411                         rotate_right(s, t);
 412                         s = m == mp->l ? mp->r : mp->l;
 413                     } else if (m == mp->r && s->c == BLACK && s->r->c == RED && s->l->c == BLACK) {
 414                         s->c = RED;
 415                         s->r->c = BLACK;
 416                         rotate_left(s, t);
 417                         s = m == mp->l ? mp->r : mp->l;
 418                     }
 419                     s->c = mp->c;
 420                     mp->c = BLACK;
 421                     if (m == mp->l) {
 422                         s->r->c = BLACK;
 423                         rotate_left(mp, t);
 424                     } else {
 425                         s->l->c = BLACK;
 426                         rotate_right(mp, t);
 427                     }
 428                 }
 429             }
 430         }
 431     }
 432     t->N--;
 433     n->k = k;                   /* n may have changed during remove */
 434     return n;                   /* the node that was deleted may be different from initial n */
 435 }
 436
 437 rb_node *rb_tree_resort(rb_tree * t, rb_node * n)
 438 {
 439     n = rb_tree_remove(t, n);
 440     insert_node(t, n);
 441     return n;
 442 }
 443
 444 /* shift all key pointers by kshift ... this is useful when the keys
 445    are pointers into another array, that has been resized with realloc */
 446 static void shift_keys(rb_node * n, ptrdiff_t kshift)
 447 {                               /* assumes n != NIL */
 448     n->k += kshift;
 449     if (n->l != NIL)
 450         shift_keys(n->l, kshift);
 451     if (n->r != NIL)
 452         shift_keys(n->r, kshift);
 453 }
 454
 455 void rb_tree_shift_keys(rb_tree * t, ptrdiff_t kshift)
 456 {
 457     if (t->root != NIL)
 458         shift_keys(t->root, kshift);
 459 }