begriffs open source - libderp/blob - src/treemap.c

   1 #include "internal/alloc.h"
   2 #include "derp/list.h"
   3 #include "derp/treemap.h"
   4
   5 /* AA (Arne Andersson) Tree:
   6  * Balanced Search Trees Made Simple
   7  * https://user.it.uu.se/~arnea/ps/simp.pdf
   8  *
   9  * As fast as an RB-tree, but free of those
  10  * ugly special cases. */
  11
  12 struct tm_node
  13 {
  14         int level;
  15         struct map_pair *pair;
  16         struct tm_node *left, *right;
  17 };
  18
  19 struct treemap
  20 {
  21         struct tm_node *root, *bottom;
  22         struct tm_node *deleted, *last;
  23
  24         dtor *key_dtor;
  25         dtor *val_dtor;
  26         comparator *cmp;
  27         void *cmp_aux;
  28         void *dtor_aux;
  29 };
  30
  31 struct tm_iter
  32 {
  33         list *stack;
  34         struct tm_node *n, *bottom;
  35 };
  36
  37 treemap *
  38 tm_new(comparator *cmp, void *cmp_aux)
  39 {
  40         treemap *t = internal_malloc(sizeof *t);
  41         struct tm_node *bottom = internal_malloc(sizeof *bottom);
  42         if (!t || !bottom)
  43         {
  44                 internal_free(t);
  45                 internal_free(bottom);
  46                 return NULL;
  47         }
  48         /* sentinel living below all leaves */
  49         *bottom = (struct tm_node){
  50                 .left = bottom, .right = bottom, .level = 0
  51         };
  52         *t = (treemap){
  53                 .root = bottom,
  54                 .bottom = bottom,
  55                 .deleted = bottom,
  56                 .last = bottom,
  57                 .cmp = cmp,
  58                 .cmp_aux = cmp_aux
  59         };
  60         return t;
  61 }
  62
  63 void
  64 tm_free(treemap *t)
  65 {
  66         if (!t)
  67                 return;
  68         tm_clear(t);
  69         internal_free(t->bottom);
  70         internal_free(t);
  71 }
  72
  73 void
  74 tm_dtor(treemap *t, dtor *key_dtor, dtor *val_dtor, void *dtor_aux)
  75 {
  76         if (!t)
  77                 return;
  78         t->key_dtor = key_dtor;
  79         t->val_dtor = val_dtor;
  80         t->dtor_aux = dtor_aux;
  81 }
  82
  83 static size_t
  84 internal_tm_length(const struct tm_node *n,
  85                    const struct tm_node *bottom)
  86 {
  87         if (n == bottom)
  88                 return 0;
  89         return 1 +
  90                 internal_tm_length(n->left, bottom) +
  91                 internal_tm_length(n->right, bottom);
  92 }
  93
  94 size_t
  95 tm_length(const treemap *t)
  96 {
  97         return t ? internal_tm_length(t->root, t->bottom) : 0;
  98 }
  99
 100 bool
 101 tm_is_empty(const treemap *t)
 102 {
 103         return tm_length(t) == 0;
 104 }
 105
 106 static void *
 107 internal_tm_at(const treemap *t, const struct tm_node *n,
 108                const void *key)
 109 {
 110         if (n == t->bottom)
 111                 return NULL;
 112         int x = t->cmp(key, n->pair->k, t->cmp_aux);
 113         if (x == 0)
 114                 return n->pair->v;
 115         else if (x < 0)
 116                 return internal_tm_at(t, n->left, key);
 117         return internal_tm_at(t, n->right, key);
 118 }
 119
 120 void *
 121 tm_at(const treemap *t, const void *key)
 122 {
 123         return t ? internal_tm_at(t, t->root, key) : NULL;
 124 }
 125
 126 static struct tm_node *
 127 internal_tm_skew(struct tm_node *n) {
 128         if (n->level != n->left->level)
 129                 return n;
 130         struct tm_node *left = n->left;
 131         n->left = left->right;
 132         left->right = n;
 133         n = left;
 134         return n;
 135 }
 136
 137 static struct tm_node *
 138 internal_tm_split(struct tm_node *n) {
 139         if(n->right->right->level != n->level)
 140                 return n;
 141         struct tm_node *right = n->right;
 142         n->right = right->left;
 143         right->left = n;
 144         n = right;
 145         n->level++;
 146         return n;
 147 }
 148
 149 static struct tm_node *
 150 internal_tm_insert(treemap *t, struct tm_node *n,
 151                    struct tm_node *prealloc)
 152 {
 153         if (n == t->bottom)
 154                 return prealloc;
 155         int x = t->cmp(prealloc->pair->k, n->pair->k, t->cmp_aux);
 156         if (x < 0)
 157                 n->left = internal_tm_insert(t, n->left, prealloc);
 158         else if (x > 0)
 159                 n->right = internal_tm_insert(t, n->right, prealloc);
 160         else
 161         {
 162                 /* prealloc was for naught, but we'll use its value */
 163                 if (n->pair->v != prealloc->pair->v && t->val_dtor)
 164                         t->val_dtor(n->pair->v, t->dtor_aux);
 165                 if (n->pair->k != prealloc->pair->k && t->key_dtor)
 166                         t->key_dtor(n->pair->k, t->dtor_aux);
 167                 *n->pair = *prealloc->pair;
 168                 internal_free(prealloc->pair);
 169                 internal_free(prealloc);
 170                 return n;
 171         }
 172         return internal_tm_split(internal_tm_skew(n));
 173 }
 174
 175 bool
 176 tm_insert(treemap *t, void *key, void *val)
 177 {
 178         if (!t)
 179                 return false;
 180         /* attempt the malloc before potentially splitting
 181          * and skewing the tree, so the insertion can be a
 182          * no-op on failure */
 183         struct tm_node *prealloc = internal_malloc(sizeof *prealloc);
 184         struct map_pair *p = internal_malloc(sizeof *p);
 185         if (!prealloc || !p)
 186         {
 187                 internal_free(prealloc);
 188                 internal_free(p);
 189                 return false;
 190         }
 191         *p = (struct map_pair){.k = key, .v = val};
 192         *prealloc = (struct tm_node){
 193                 .level = 1, .pair = p, .left = t->bottom, .right = t->bottom
 194         };
 195
 196         t->root = internal_tm_insert(t, t->root, prealloc);
 197         return true;
 198 }
 199
 200 static struct tm_node *
 201 internal_tm_remove(treemap *t, struct tm_node *n, void *key)
 202 {
 203         if (n == t->bottom)
 204                 return n;
 205
 206         /* 1: search down the tree and set pointers last and deleted */
 207
 208         t->last = n;
 209         if (t->cmp(key, n->pair->k, t->cmp_aux) < 0)
 210                 n->left = internal_tm_remove(t, n->left, key);
 211         else
 212         {
 213                 t->deleted = n;
 214                 n->right = internal_tm_remove(t, n->right, key);
 215         }
 216
 217         /* 2: At the bottom of the tree, remove element if present */
 218
 219         if (n == t->last && t->deleted != t->bottom &&
 220             t->cmp(key, t->deleted->pair->k, t->cmp_aux) == 0)
 221         {
 222                 if (t->key_dtor)
 223                         t->key_dtor(t->deleted->pair->k, t->dtor_aux);
 224                 if (t->val_dtor)
 225                         t->val_dtor(t->deleted->pair->v, t->dtor_aux);
 226
 227                 *t->deleted->pair = *n->pair;
 228                 t->deleted = t->bottom;
 229                 n = n->right;
 230
 231                 internal_free(t->last->pair);
 232                 internal_free(t->last);
 233         } /* 3: on the way back up, rebalance */
 234         else if (n->left->level  < n->level-1 ||
 235                  n->right->level < n->level-1) {
 236                 n->level--;
 237                 if (n->right->level > n->level)
 238                         n->right->level = n->level;
 239                 n               = internal_tm_skew(n);
 240                 n->right        = internal_tm_skew(n->right);
 241                 n->right->right = internal_tm_skew(n->right->right);
 242                 n               = internal_tm_split(n);
 243                 n->right        = internal_tm_split(n->right);
 244         }
 245         return n;
 246 }
 247
 248 bool
 249 tm_remove(treemap *t, void *key)
 250 {
 251         if (!t)
 252                 return false;
 253         t->root = internal_tm_remove(t, t->root, key);
 254         return true; // TODO: return false if key wasn't found
 255 }
 256
 257 static void
 258 internal_tm_clear(treemap *t, struct tm_node *n)
 259 {
 260         if (n == t->bottom)
 261                 return;
 262         internal_tm_clear(t, n->left);
 263         internal_tm_clear(t, n->right);
 264         if (t->key_dtor)
 265                 t->key_dtor(n->pair->k, t->dtor_aux);
 266         if (t->val_dtor)
 267                 t->val_dtor(n->pair->v, t->dtor_aux);
 268         internal_free(n->pair);
 269         internal_free(n);
 270 }
 271
 272 void
 273 tm_clear(treemap *t)
 274 {
 275         if (!t)
 276                 return;
 277         internal_tm_clear(t, t->root);
 278         t->root = t->deleted = t->last = t->bottom;
 279 }
 280
 281 tm_iter *
 282 tm_iter_begin(treemap *t)
 283 {
 284         if (!t)
 285                 return NULL;
 286         struct tm_iter *i = internal_malloc(sizeof *i);
 287         list *l = l_new();
 288         if (!i || !l)
 289         {
 290                 internal_free(i);
 291                 l_free(l);
 292                 return NULL;
 293         }
 294         *i = (struct tm_iter){
 295                 .stack = l,
 296                 .n = t->root,
 297                 .bottom = t->bottom
 298         };
 299         return i;
 300 }
 301
 302 struct map_pair *
 303 tm_iter_next(tm_iter *i)
 304 {
 305         if (!i)
 306                 return NULL;
 307         if (l_is_empty(i->stack) && i->n == i->bottom)
 308                 return NULL; /* done */
 309         if (i->n != i->bottom)
 310         {
 311                 if (!l_append(i->stack, i->n))
 312                         return NULL; /* OOM */
 313                 i->n = i->n->left;
 314                 return tm_iter_next(i);
 315         }
 316         struct tm_node *result = l_remove_last(i->stack);
 317         i->n = result->right;
 318         return result->pair;
 319 }
 320
 321 void
 322 tm_iter_free(tm_iter *i)
 323 {
 324         if (i)
 325                 l_free(i->stack);
 326         internal_free(i);
 327 }