Dogcows Code - chaz/rasterize/blob - rbtree.c

   1
   2 /*
   3  * CS5600 University of Utah
   4  * Charles McGarvey
   5  * mcgarvey@eng.utah.edu
   6  */
   7
   8 #include "rbtree.h"
   9
  10
  11 #define BLACK       (RBTREE_BLACK)
  12 #define RED         (RBTREE_RED)
  13 #define SENTINEL    (rbtree_sentinel)
  14
  15
  16 rbtree_node_t _rbtree_sentinel = {
  17     &_rbtree_sentinel, &_rbtree_sentinel, &_rbtree_sentinel,
  18     BLACK
  19 };
  20 rbtree_node_t* rbtree_sentinel = &_rbtree_sentinel;
  21
  22
  23 /*
  24  * Allocate and initialize a new node for the tree.
  25  */
  26 INLINE_MAYBE
  27 rbtree_node_t* rbtree_node_alloc(rbtree_t* t)
  28 {
  29     rbtree_node_t* n = mem_alloc(sizeof(rbtree_node_t) + t->size);
  30     n->parent = NULL;
  31     n->left = n->right = rbtree_sentinel;
  32     n->color = RBTREE_BLACK;
  33     return n;
  34 }
  35
  36
  37 INLINE_MAYBE
  38 void rbtree_init(rbtree_t* t, size_t size, rbtree_search_fn_t search)
  39 {
  40     t->root = rbtree_sentinel;
  41     t->size = size;
  42     t->search = search;
  43 }
  44
  45 rbtree_t* rbtree_alloc(size_t size, rbtree_search_fn_t search)
  46 {
  47     rbtree_t* t = mem_alloc(sizeof(rbtree_t));
  48     rbtree_init(t, size, search);
  49     return t;
  50 }
  51
  52
  53 void rbtree_destroy(rbtree_t* t)
  54 {
  55     rbtree_clear(t);
  56     mem_free(t);
  57 }
  58
  59
  60 INLINE_MAYBE
  61 void _rotate_left(rbtree_node_t** r, rbtree_node_t* n)
  62 {
  63     rbtree_node_t* t = n->right;
  64     n->right = t->left;
  65
  66     if (t->left != SENTINEL) {
  67         t->left->parent = n;
  68     }
  69     t->parent = n->parent;
  70
  71     if (n == *r) {
  72         *r = t;
  73     }
  74     else if (n == n->parent->left) {
  75         n->parent->left = t;
  76     }
  77     else {
  78         n->parent->right = t;
  79     }
  80
  81     t->left = n;
  82     n->parent = t;
  83 }
  84
  85 INLINE_MAYBE
  86 void _rotate_right(rbtree_node_t** r, rbtree_node_t* n)
  87 {
  88     rbtree_node_t* t = n->left;
  89     n->left = t->right;
  90
  91     if (t->right != SENTINEL) {
  92         t->right->parent = n;
  93     }
  94     t->parent = n->parent;
  95
  96     if (n == *r) {
  97         *r = t;
  98     }
  99     else if (n == n->parent->right) {
 100         n->parent->right = t;
 101     }
 102     else {
 103         n->parent->left = t;
 104     }
 105
 106     t->right = n;
 107     n->parent = t;
 108 }
 109
 110
 111 rbtree_node_t* rbtree_insert(rbtree_t* t, void* data)
 112 {
 113     rbtree_node_t* r = t->root;
 114     rbtree_node_t* n = r;
 115
 116     int c = t->search(&n, data);
 117     if (c == 0) {
 118         memcpy(RBTREE_NODE_DATA(n), data, t->size);
 119         return n;
 120     }
 121     else if (n == SENTINEL) {
 122         r = rbtree_node_alloc(t);
 123         memcpy(RBTREE_NODE_DATA(r), data, t->size);
 124         t->root = r;
 125         return r;
 126     }
 127
 128     rbtree_node_t* m = rbtree_node_alloc(t);
 129     m->parent = n;
 130     m->color = RED;
 131     memcpy(RBTREE_NODE_DATA(m), data, t->size);
 132     if (c < 0) {
 133         n->left = m;
 134     }
 135     else {
 136         n->right = m;
 137     }
 138
 139     n = m;
 140     while (n != r && n->parent->color == RED) {
 141         if (n->parent == n->parent->parent->left) {
 142             rbtree_node_t* t = n->parent->parent->right;
 143             if (t->color == RED) {
 144                 n->parent->color = BLACK;
 145                 t->color = BLACK;
 146                 n->parent->parent->color = RED;
 147                 n = n->parent->parent;
 148             }
 149             else {
 150                 if (n == n->parent->right) {
 151                     n = n->parent;
 152                     _rotate_left(&r, n);
 153                 }
 154                 n->parent->color = BLACK;
 155                 n->parent->parent->color = RED;
 156                 _rotate_right(&r, n->parent->parent);
 157             }
 158         }
 159         else {
 160             rbtree_node_t* t = n->parent->parent->left;
 161             if (t->color == RED) {
 162                 n->parent->color = BLACK;
 163                 t->color = BLACK;
 164                 n->parent->parent->color = RED;
 165                 n = n->parent->parent;
 166             }
 167             else {
 168                 if (n == n->parent->left) {
 169                     n = n->parent;
 170                     _rotate_right(&r, n);
 171                 }
 172                 n->parent->color = BLACK;
 173                 n->parent->parent->color = RED;
 174                 _rotate_left(&r, n->parent->parent);
 175             }
 176         }
 177     }
 178     t->root = r;
 179     r->color = BLACK;
 180
 181     return m;
 182 }
 183
 184 rbtree_node_t* rbtree_delete(rbtree_t* t, void* data)
 185 {
 186     rbtree_node_t* n = t->root;
 187     if (t->search(&n, data) == 0) {
 188         return rbtree_delete_node(t, n);
 189     }
 190     return NULL;
 191 }
 192
 193 rbtree_node_t* rbtree_search(rbtree_t* t, void* data)
 194 {
 195     rbtree_node_t* n = t->root;
 196     if (t->search(&n, data) == 0) {
 197         return n;
 198     }
 199     return NULL;
 200 }
 201
 202
 203 rbtree_node_t* rbtree_delete_node(rbtree_t* t, rbtree_node_t* n)
 204 {
 205     rbtree_node_t* r = t->root;
 206     rbtree_node_t* u;
 207     rbtree_node_t* v;
 208
 209     if (n->left == SENTINEL) {
 210         u = n->right;
 211         v = n;
 212     }
 213     else if (n->right == SENTINEL) {
 214         u = n->left;
 215         v = n;
 216     }
 217     else {
 218         v = rbtree_node_min(n->right);
 219         if (v->left != SENTINEL) {
 220             u = v->left;
 221         }
 222         else {
 223             u = v->right;
 224         }
 225     }
 226
 227     if (v == r) {
 228         r = u;
 229         u->color = BLACK;
 230         goto done;
 231     }
 232
 233     int isred = (v->color == RED);
 234
 235     if (v == v->parent->left) {
 236         v->parent->left = u;
 237     }
 238     else {
 239         v->parent->right = u;
 240     }
 241
 242     if (v == n) {
 243         u->parent = v->parent;
 244     }
 245     else {
 246         if (v->parent == n) {
 247             u->parent = v;
 248         }
 249         else {
 250             u->parent = v->parent;
 251         }
 252
 253         *v = *n;
 254
 255         if (n == r) {
 256             r = v;
 257         }
 258         else if (n == n->parent->left) {
 259             n->parent->left = v;
 260         }
 261         else {
 262             n->parent->right = v;
 263         }
 264
 265         if (v->left != SENTINEL) {
 266             v->left->parent = v;
 267         }
 268         if (v->right != SENTINEL) {
 269             v->right->parent = v;
 270         }
 271     }
 272
 273     if (isred) {
 274         goto done;
 275     }
 276
 277     while (u != r && u->color == BLACK) {
 278         if (u == u->parent->left) {
 279             rbtree_node_t* w = u->parent->right;
 280             if (w->color == RED) {
 281                 w->color = BLACK;
 282                 w->parent->color = RED;
 283                 _rotate_left(&r, u->parent);
 284                 w = u->parent->right;
 285             }
 286             if (w->left->color == BLACK && w->right->color == BLACK) {
 287                 w->color = RED;
 288                 u = u->parent;
 289             }
 290             else {
 291                 if (w->right->color == BLACK) {
 292                     w->left->color = BLACK;
 293                     w->color = RED;
 294                     _rotate_right(&r, w);
 295                     w = u->parent->right;
 296                 }
 297                 w->color = u->parent->color;
 298                 u->parent->color = BLACK;
 299                 w->right->color = BLACK;
 300                 _rotate_left(&r, u->parent);
 301                 u = r;
 302             }
 303         }
 304         else {
 305             rbtree_node_t* w = u->parent->left;
 306             if (w->color == RED) {
 307                 w->color = BLACK;
 308                 w->parent->color = RED;
 309                 _rotate_right(&r, u->parent);
 310                 w = u->parent->left;
 311             }
 312             if (w->left->color == BLACK && w->right->color == BLACK) {
 313                 w->color = RED;
 314                 u = u->parent;
 315             }
 316             else {
 317                 if (w->left->color == BLACK) {
 318                     w->right->color = BLACK;
 319                     w->color = RED;
 320                     _rotate_left(&r, w);
 321                     w = u->parent->left;
 322                 }
 323                 w->color = u->parent->color;
 324                 u->parent->color = BLACK;
 325                 w->left->color = BLACK;
 326                 _rotate_right(&r, u->parent);
 327                 u = r;
 328             }
 329         }
 330     }
 331
 332 done:
 333     t->root = r;
 334     n->parent = n->left = n->right = NULL;
 335     return n;
 336 }
 337
 338
 339 void node_call(rbtree_node_t* n, rbtree_call_fn_t fn)
 340 {
 341     if (n == SENTINEL) {
 342         return;
 343     }
 344     node_call(n->left, fn);
 345     fn(RBTREE_NODE_DATA(n));
 346     node_call(n->right, fn);
 347 }
 348
 349 void rbtree_call(const rbtree_t* t, rbtree_call_fn_t fn)
 350 {
 351     node_call(t->root, fn);
 352 }
 353
 354
 355 /*
 356  * Prune off a branch of the tree.
 357  */
 358 static void _node_prune(rbtree_node_t* n)
 359 {
 360     if (n == SENTINEL) {
 361         return;
 362     }
 363     _node_prune(n->left);
 364     _node_prune(n->right);
 365     mem_free(n);
 366 }
 367
 368 void rbtree_clear(rbtree_t* t)
 369 {
 370     if (t == NULL) {
 371         return;
 372     }
 373     _node_prune(t->root);
 374 }
 375