map.c

   1 /* Crit-bit tree based map which supports lookups based on unique
   2  * prefixes as well as ordered iteration.
   3  *
   4  * Based on public domain code from Rusty Russel, Adam Langley
   5  * and D. J. Bernstein.
   6  *
   7  * Further information about the data structure can be found at:
   8  *  http://cr.yp.to/critbit.html
   9  *  http://github.com/agl/critbit
  10  *  http://ccodearchive.net/info/strmap.html
  11  */
  12 #include <stdlib.h>
  13 #include <string.h>
  14 #include <errno.h>
  15 #include <inttypes.h>
  16 #include "map.h"
  17
  18 typedef struct Node Node;
  19
  20 struct Map {     /* struct holding either an item with value v and key u.s or an internal node */
  21         union {
  22                 Node *n;
  23                 const char *s;
  24         } u;
  25         void *v; /* value stored in the map, if non NULL u.s holds the corresponding key */
  26 };
  27
  28 struct Node {
  29         Map child[2];    /* These point to strings or nodes. */
  30         size_t byte_num; /* The byte number where first bit differs. */
  31         uint8_t bit_num; /* The bit where these children differ. */
  32 };
  33
  34 /* Closest key to this in a non-empty map. */
  35 static Map *closest(Map *n, const char *key)
  36 {
  37         size_t len = strlen(key);
  38         const uint8_t *bytes = (const uint8_t *)key;
  39
  40         /* Anything with NULL value is an internal node. */
  41         while (!n->v) {
  42                 uint8_t direction = 0;
  43
  44                 if (n->u.n->byte_num < len) {
  45                         uint8_t c = bytes[n->u.n->byte_num];
  46                         direction = (c >> n->u.n->bit_num) & 1;
  47                 }
  48                 n = &n->u.n->child[direction];
  49         }
  50         return n;
  51 }
  52
  53 void *map_get(const Map *map, const char *key)
  54 {
  55         /* Not empty map? */
  56         if (map->u.n) {
  57                 Map *n = closest((Map *)map, key);
  58                 if (strcmp(key, n->u.s) == 0)
  59                         return n->v;
  60         }
  61         return NULL;
  62 }
  63
  64 void *map_closest(const Map *map, const char *prefix)
  65 {
  66         errno = 0;
  67         void *v = map_get(map, prefix);
  68         if (v)
  69                 return v;
  70         const Map *m = map_prefix(map, prefix);
  71         if (map_empty(m))
  72                 errno = ENOENT;
  73         return m->v;
  74 }
  75
  76 bool map_contains(const Map *map, const char *prefix)
  77 {
  78         return !map_empty(map_prefix(map, prefix));
  79 }
  80
  81 static bool leaf(const char *key, void *value, void *data)
  82 {
  83         int *nodes = data;
  84         return (*nodes)++ < 1;
  85 }
  86
  87 bool map_leaf(const Map *map, const char *prefix)
  88 {
  89         int nodes = 0;
  90         map_iterate(map_prefix(map, prefix), leaf, &nodes);
  91         return nodes == 1;
  92 }
  93
  94 bool map_put(Map *map, const char *k, const void *value)
  95 {
  96         size_t len = strlen(k);
  97         const uint8_t *bytes = (const uint8_t *)k;
  98         Map *n;
  99         Node *newn;
 100         size_t byte_num;
 101         uint8_t bit_num, new_dir;
 102         char *key;
 103
 104         if (!value) {
 105                 errno = EINVAL;
 106                 return false;
 107         }
 108
 109         if (!(key = strdup(k))) {
 110                 errno = ENOMEM;
 111                 return false;
 112         }
 113
 114         /* Empty map? */
 115         if (!map->u.n) {
 116                 map->u.s = key;
 117                 map->v = (void *)value;
 118                 return true;
 119         }
 120
 121         /* Find closest existing key. */
 122         n = closest(map, key);
 123
 124         /* Find where they differ. */
 125         for (byte_num = 0; n->u.s[byte_num] == key[byte_num]; byte_num++) {
 126                 if (key[byte_num] == '\0') {
 127                         /* All identical! */
 128                         free(key);
 129                         errno = EEXIST;
 130                         return false;
 131                 }
 132         }
 133
 134         /* Find which bit differs */
 135         uint8_t diff = (uint8_t)n->u.s[byte_num] ^ bytes[byte_num];
 136         /* TODO: bit_num = 31 - __builtin_clz(diff); ? */
 137         for (bit_num = 0; diff >>= 1; bit_num++);
 138
 139         /* Which direction do we go at this bit? */
 140         new_dir = ((bytes[byte_num]) >> bit_num) & 1;
 141
 142         /* Allocate new node. */
 143         newn = malloc(sizeof(*newn));
 144         if (!newn) {
 145                 free(key);
 146                 errno = ENOMEM;
 147                 return false;
 148         }
 149         newn->byte_num = byte_num;
 150         newn->bit_num = bit_num;
 151         newn->child[new_dir].v = (void *)value;
 152         newn->child[new_dir].u.s = key;
 153
 154         /* Find where to insert: not closest, but first which differs! */
 155         n = map;
 156         while (!n->v) {
 157                 uint8_t direction = 0;
 158
 159                 if (n->u.n->byte_num > byte_num)
 160                         break;
 161                 /* Subtle: bit numbers are "backwards" for comparison */
 162                 if (n->u.n->byte_num == byte_num && n->u.n->bit_num < bit_num)
 163                         break;
 164
 165                 if (n->u.n->byte_num < len) {
 166                         uint8_t c = bytes[n->u.n->byte_num];
 167                         direction = (c >> n->u.n->bit_num) & 1;
 168                 }
 169                 n = &n->u.n->child[direction];
 170         }
 171
 172         newn->child[!new_dir] = *n;
 173         n->u.n = newn;
 174         n->v = NULL;
 175         return true;
 176 }
 177
 178 void *map_delete(Map *map, const char *key)
 179 {
 180         size_t len = strlen(key);
 181         const uint8_t *bytes = (const uint8_t *)key;
 182         Map *parent = NULL, *n;
 183         void *value = NULL;
 184         uint8_t direction;
 185
 186         /* Empty map? */
 187         if (!map->u.n) {
 188                 errno = ENOENT;
 189                 return NULL;
 190         }
 191
 192         /* Find closest, but keep track of parent. */
 193         n = map;
 194         /* Anything with NULL value is a node. */
 195         while (!n->v) {
 196                 uint8_t c = 0;
 197
 198                 parent = n;
 199                 if (n->u.n->byte_num < len) {
 200                         c = bytes[n->u.n->byte_num];
 201                         direction = (c >> n->u.n->bit_num) & 1;
 202                 } else {
 203                         direction = 0;
 204                 }
 205                 n = &n->u.n->child[direction];
 206         }
 207
 208         /* Did we find it? */
 209         if (strcmp(key, n->u.s)) {
 210                 errno = ENOENT;
 211                 return NULL;
 212         }
 213
 214         free((char*)n->u.s);
 215         value = n->v;
 216
 217         if (!parent) {
 218                 /* We deleted last node. */
 219                 map->u.n = NULL;
 220         } else {
 221                 Node *old = parent->u.n;
 222                 /* Raise other node to parent. */
 223                 *parent = old->child[!direction];
 224                 free(old);
 225         }
 226
 227         return value;
 228 }
 229
 230 static bool iterate(Map n, bool (*handle)(const char *, void *, void *), const void *data)
 231 {
 232         if (n.v)
 233                 return handle(n.u.s, n.v, (void *)data);
 234
 235         return iterate(n.u.n->child[0], handle, data)
 236                 && iterate(n.u.n->child[1], handle, data);
 237 }
 238
 239 void map_iterate(const Map *map, bool (*handle)(const char *, void *, void *), const void *data)
 240 {
 241         /* Empty map? */
 242         if (!map->u.n)
 243                 return;
 244
 245         iterate(*map, handle, data);
 246 }
 247
 248 typedef struct {
 249         const char *key;
 250         void *value;
 251 } KeyValue;
 252
 253 static bool first(const char *key, void *value, void *data)
 254 {
 255         KeyValue *kv = data;
 256         kv->key = key;
 257         kv->value = value;
 258         return false;
 259 }
 260
 261 void *map_first(const Map *map, const char **key)
 262 {
 263         KeyValue kv = { 0 };
 264         map_iterate(map, first, &kv);
 265         if (key && kv.key)
 266                 *key = kv.key;
 267         return kv.value;
 268 }
 269
 270 const Map *map_prefix(const Map *map, const char *prefix)
 271 {
 272         const Map *n, *top;
 273         size_t len = strlen(prefix);
 274         const uint8_t *bytes = (const uint8_t *)prefix;
 275
 276         /* Empty map -> return empty map. */
 277         if (!map->u.n)
 278                 return map;
 279
 280         top = n = map;
 281
 282         /* We walk to find the top, but keep going to check prefix matches. */
 283         while (!n->v) {
 284                 uint8_t c = 0, direction;
 285
 286                 if (n->u.n->byte_num < len)
 287                         c = bytes[n->u.n->byte_num];
 288
 289                 direction = (c >> n->u.n->bit_num) & 1;
 290                 n = &n->u.n->child[direction];
 291                 if (c)
 292                         top = n;
 293         }
 294
 295         if (strncmp(n->u.s, prefix, len)) {
 296                 /* Convenient return for prefixes which do not appear in map. */
 297                 static const Map empty_map;
 298                 return &empty_map;
 299         }
 300
 301         return top;
 302 }
 303
 304 static void clear(Map n)
 305 {
 306         if (!n.v) {
 307                 clear(n.u.n->child[0]);
 308                 clear(n.u.n->child[1]);
 309                 free(n.u.n);
 310         } else {
 311                 free((char*)n.u.s);
 312         }
 313 }
 314
 315 void map_clear(Map *map)
 316 {
 317         if (map->u.n)
 318                 clear(*map);
 319         map->u.n = NULL;
 320         map->v = NULL;
 321 }
 322
 323 static bool copy(Map *dest, Map n)
 324 {
 325         if (!n.v) {
 326                 return copy(dest, n.u.n->child[0]) &&
 327                        copy(dest, n.u.n->child[1]);
 328         } else {
 329                 if (!map_put(dest, n.u.s, n.v) && map_get(dest, n.u.s) != n.v) {
 330                         map_delete(dest, n.u.s);
 331                         return map_put(dest, n.u.s, n.v);
 332                 }
 333                 return true;
 334         }
 335 }
 336
 337 bool map_copy(Map *dest, Map *src)
 338 {
 339         if (!src || !src->u.n)
 340                 return true;
 341
 342         return copy(dest, *src);
 343 }
 344
 345 bool map_empty(const Map *map)
 346 {
 347         return map->u.n == NULL;
 348 }
 349
 350 Map *map_new(void)
 351 {
 352         return calloc(1, sizeof(Map));
 353 }
 354
 355 void map_free(Map *map)
 356 {
 357         if (!map)
 358                 return;
 359         map_clear(map);
 360         free(map);
 361 }
 362
 363 static bool free_elem(const char *key, void *value, void *data)
 364 {
 365         free(value);
 366         return true;
 367 }
 368
 369 void map_free_full(Map *map)
 370 {
 371         if (!map)
 372                 return;
 373         map_iterate(map, free_elem, NULL);
 374         map_free(map);
 375 }