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 bool map_put(Map *map, const char *k, const void *value)
  82 {
  83         size_t len = strlen(k);
  84         const uint8_t *bytes = (const uint8_t *)k;
  85         Map *n;
  86         Node *newn;
  87         size_t byte_num;
  88         uint8_t bit_num, new_dir;
  89         char *key;
  90
  91         if (!value) {
  92                 errno = EINVAL;
  93                 return false;
  94         }
  95
  96         if (!(key = strdup(k))) {
  97                 errno = ENOMEM;
  98                 return false;
  99         }
 100
 101         /* Empty map? */
 102         if (!map->u.n) {
 103                 map->u.s = key;
 104                 map->v = (void *)value;
 105                 return true;
 106         }
 107
 108         /* Find closest existing key. */
 109         n = closest(map, key);
 110
 111         /* Find where they differ. */
 112         for (byte_num = 0; n->u.s[byte_num] == key[byte_num]; byte_num++) {
 113                 if (key[byte_num] == '\0') {
 114                         /* All identical! */
 115                         free(key);
 116                         errno = EEXIST;
 117                         return false;
 118                 }
 119         }
 120
 121         /* Find which bit differs */
 122         uint8_t diff = (uint8_t)n->u.s[byte_num] ^ bytes[byte_num];
 123         /* TODO: bit_num = 31 - __builtin_clz(diff); ? */
 124         for (bit_num = 0; diff >>= 1; bit_num++);
 125
 126         /* Which direction do we go at this bit? */
 127         new_dir = ((bytes[byte_num]) >> bit_num) & 1;
 128
 129         /* Allocate new node. */
 130         newn = malloc(sizeof(*newn));
 131         if (!newn) {
 132                 free(key);
 133                 errno = ENOMEM;
 134                 return false;
 135         }
 136         newn->byte_num = byte_num;
 137         newn->bit_num = bit_num;
 138         newn->child[new_dir].v = (void *)value;
 139         newn->child[new_dir].u.s = key;
 140
 141         /* Find where to insert: not closest, but first which differs! */
 142         n = map;
 143         while (!n->v) {
 144                 uint8_t direction = 0;
 145
 146                 if (n->u.n->byte_num > byte_num)
 147                         break;
 148                 /* Subtle: bit numbers are "backwards" for comparison */
 149                 if (n->u.n->byte_num == byte_num && n->u.n->bit_num < bit_num)
 150                         break;
 151
 152                 if (n->u.n->byte_num < len) {
 153                         uint8_t c = bytes[n->u.n->byte_num];
 154                         direction = (c >> n->u.n->bit_num) & 1;
 155                 }
 156                 n = &n->u.n->child[direction];
 157         }
 158
 159         newn->child[!new_dir] = *n;
 160         n->u.n = newn;
 161         n->v = NULL;
 162         return true;
 163 }
 164
 165 void *map_delete(Map *map, const char *key)
 166 {
 167         size_t len = strlen(key);
 168         const uint8_t *bytes = (const uint8_t *)key;
 169         Map *parent = NULL, *n;
 170         void *value = NULL;
 171         uint8_t direction;
 172
 173         /* Empty map? */
 174         if (!map->u.n) {
 175                 errno = ENOENT;
 176                 return NULL;
 177         }
 178
 179         /* Find closest, but keep track of parent. */
 180         n = map;
 181         /* Anything with NULL value is a node. */
 182         while (!n->v) {
 183                 uint8_t c = 0;
 184
 185                 parent = n;
 186                 if (n->u.n->byte_num < len) {
 187                         c = bytes[n->u.n->byte_num];
 188                         direction = (c >> n->u.n->bit_num) & 1;
 189                 } else {
 190                         direction = 0;
 191                 }
 192                 n = &n->u.n->child[direction];
 193         }
 194
 195         /* Did we find it? */
 196         if (strcmp(key, n->u.s)) {
 197                 errno = ENOENT;
 198                 return NULL;
 199         }
 200
 201         free((char*)n->u.s);
 202         value = n->v;
 203
 204         if (!parent) {
 205                 /* We deleted last node. */
 206                 map->u.n = NULL;
 207         } else {
 208                 Node *old = parent->u.n;
 209                 /* Raise other node to parent. */
 210                 *parent = old->child[!direction];
 211                 free(old);
 212         }
 213
 214         return value;
 215 }
 216
 217 static bool iterate(Map n, bool (*handle)(const char *, void *, void *), const void *data)
 218 {
 219         if (n.v)
 220                 return handle(n.u.s, n.v, (void *)data);
 221
 222         return iterate(n.u.n->child[0], handle, data)
 223                 && iterate(n.u.n->child[1], handle, data);
 224 }
 225
 226 void map_iterate(const Map *map, bool (*handle)(const char *, void *, void *), const void *data)
 227 {
 228         /* Empty map? */
 229         if (!map->u.n)
 230                 return;
 231
 232         iterate(*map, handle, data);
 233 }
 234
 235 typedef struct {
 236         const char *key;
 237         void *value;
 238 } KeyValue;
 239
 240 static bool first(const char *key, void *value, void *data)
 241 {
 242         KeyValue *kv = data;
 243         kv->key = key;
 244         kv->value = value;
 245         return false;
 246 }
 247
 248 void *map_first(const Map *map, const char **key)
 249 {
 250         KeyValue kv = { 0 };
 251         map_iterate(map, first, &kv);
 252         if (key && kv.key)
 253                 *key = kv.key;
 254         return kv.value;
 255 }
 256
 257 const Map *map_prefix(const Map *map, const char *prefix)
 258 {
 259         const Map *n, *top;
 260         size_t len = strlen(prefix);
 261         const uint8_t *bytes = (const uint8_t *)prefix;
 262
 263         /* Empty map -> return empty map. */
 264         if (!map->u.n)
 265                 return map;
 266
 267         top = n = map;
 268
 269         /* We walk to find the top, but keep going to check prefix matches. */
 270         while (!n->v) {
 271                 uint8_t c = 0, direction;
 272
 273                 if (n->u.n->byte_num < len)
 274                         c = bytes[n->u.n->byte_num];
 275
 276                 direction = (c >> n->u.n->bit_num) & 1;
 277                 n = &n->u.n->child[direction];
 278                 if (c)
 279                         top = n;
 280         }
 281
 282         if (strncmp(n->u.s, prefix, len)) {
 283                 /* Convenient return for prefixes which do not appear in map. */
 284                 static const Map empty_map;
 285                 return &empty_map;
 286         }
 287
 288         return top;
 289 }
 290
 291 static void clear(Map n)
 292 {
 293         if (!n.v) {
 294                 clear(n.u.n->child[0]);
 295                 clear(n.u.n->child[1]);
 296                 free(n.u.n);
 297         } else {
 298                 free((char*)n.u.s);
 299         }
 300 }
 301
 302 void map_clear(Map *map)
 303 {
 304         if (map->u.n)
 305                 clear(*map);
 306         map->u.n = NULL;
 307         map->v = NULL;
 308 }
 309
 310 static bool copy(Map *dest, Map n)
 311 {
 312         if (!n.v) {
 313                 return copy(dest, n.u.n->child[0]) &&
 314                        copy(dest, n.u.n->child[1]);
 315         } else {
 316                 if (!map_put(dest, n.u.s, n.v) && map_get(dest, n.u.s) != n.v) {
 317                         map_delete(dest, n.u.s);
 318                         return map_put(dest, n.u.s, n.v);
 319                 }
 320                 return true;
 321         }
 322 }
 323
 324 bool map_copy(Map *dest, Map *src)
 325 {
 326         if (!src || !src->u.n)
 327                 return true;
 328
 329         return copy(dest, *src);
 330 }
 331
 332 bool map_empty(const Map *map)
 333 {
 334         return map->u.n == NULL;
 335 }
 336
 337 Map *map_new(void)
 338 {
 339         return calloc(1, sizeof(Map));
 340 }
 341
 342 void map_free(Map *map)
 343 {
 344         if (!map)
 345                 return;
 346         map_clear(map);
 347         free(map);
 348 }
 349
 350 static bool free_elem(const char *key, void *value, void *data)
 351 {
 352         free(value);
 353         return true;
 354 }
 355
 356 void map_free_full(Map *map)
 357 {
 358         if (!map)
 359                 return;
 360         map_iterate(map, free_elem, NULL);
 361         map_free(map);
 362 }