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 (!m->v)
  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                         return false;
 117                 }
 118         }
 119
 120         /* Find which bit differs */
 121         uint8_t diff = (uint8_t)n->u.s[byte_num] ^ bytes[byte_num];
 122         /* TODO: bit_num = 31 - __builtin_clz(diff); ? */
 123         for (bit_num = 0; diff >>= 1; bit_num++);
 124
 125         /* Which direction do we go at this bit? */
 126         new_dir = ((bytes[byte_num]) >> bit_num) & 1;
 127
 128         /* Allocate new node. */
 129         newn = malloc(sizeof(*newn));
 130         if (!newn) {
 131                 free(key);
 132                 errno = ENOMEM;
 133                 return false;
 134         }
 135         newn->byte_num = byte_num;
 136         newn->bit_num = bit_num;
 137         newn->child[new_dir].v = (void *)value;
 138         newn->child[new_dir].u.s = key;
 139
 140         /* Find where to insert: not closest, but first which differs! */
 141         n = map;
 142         while (!n->v) {
 143                 uint8_t direction = 0;
 144
 145                 if (n->u.n->byte_num > byte_num)
 146                         break;
 147                 /* Subtle: bit numbers are "backwards" for comparison */
 148                 if (n->u.n->byte_num == byte_num && n->u.n->bit_num < bit_num)
 149                         break;
 150
 151                 if (n->u.n->byte_num < len) {
 152                         uint8_t c = bytes[n->u.n->byte_num];
 153                         direction = (c >> n->u.n->bit_num) & 1;
 154                 }
 155                 n = &n->u.n->child[direction];
 156         }
 157
 158         newn->child[!new_dir] = *n;
 159         n->u.n = newn;
 160         n->v = NULL;
 161         return true;
 162 }
 163
 164 void *map_delete(Map *map, const char *key)
 165 {
 166         size_t len = strlen(key);
 167         const uint8_t *bytes = (const uint8_t *)key;
 168         Map *parent = NULL, *n;
 169         void *value = NULL;
 170         uint8_t direction;
 171
 172         /* Empty map? */
 173         if (!map->u.n) {
 174                 errno = ENOENT;
 175                 return NULL;
 176         }
 177
 178         /* Find closest, but keep track of parent. */
 179         n = map;
 180         /* Anything with NULL value is a node. */
 181         while (!n->v) {
 182                 uint8_t c = 0;
 183
 184                 parent = n;
 185                 if (n->u.n->byte_num < len) {
 186                         c = bytes[n->u.n->byte_num];
 187                         direction = (c >> n->u.n->bit_num) & 1;
 188                 } else {
 189                         direction = 0;
 190                 }
 191                 n = &n->u.n->child[direction];
 192         }
 193
 194         /* Did we find it? */
 195         if (strcmp(key, n->u.s)) {
 196                 errno = ENOENT;
 197                 return NULL;
 198         }
 199
 200         free((char*)n->u.s);
 201         value = n->v;
 202
 203         if (!parent) {
 204                 /* We deleted last node. */
 205                 map->u.n = NULL;
 206         } else {
 207                 Node *old = parent->u.n;
 208                 /* Raise other node to parent. */
 209                 *parent = old->child[!direction];
 210                 free(old);
 211         }
 212
 213         return value;
 214 }
 215
 216 static bool iterate(Map n, bool (*handle)(const char *, void *, void *), const void *data)
 217 {
 218         if (n.v)
 219                 return handle(n.u.s, n.v, (void *)data);
 220
 221         return iterate(n.u.n->child[0], handle, data)
 222                 && iterate(n.u.n->child[1], handle, data);
 223 }
 224
 225 void map_iterate(const Map *map, bool (*handle)(const char *, void *, void *), const void *data)
 226 {
 227         /* Empty map? */
 228         if (!map->u.n)
 229                 return;
 230
 231         iterate(*map, handle, data);
 232 }
 233
 234 const Map *map_prefix(const Map *map, const char *prefix)
 235 {
 236         const Map *n, *top;
 237         size_t len = strlen(prefix);
 238         const uint8_t *bytes = (const uint8_t *)prefix;
 239
 240         /* Empty map -> return empty map. */
 241         if (!map->u.n)
 242                 return map;
 243
 244         top = n = map;
 245
 246         /* We walk to find the top, but keep going to check prefix matches. */
 247         while (!n->v) {
 248                 uint8_t c = 0, direction;
 249
 250                 if (n->u.n->byte_num < len)
 251                         c = bytes[n->u.n->byte_num];
 252
 253                 direction = (c >> n->u.n->bit_num) & 1;
 254                 n = &n->u.n->child[direction];
 255                 if (c)
 256                         top = n;
 257         }
 258
 259         if (strncmp(n->u.s, prefix, len)) {
 260                 /* Convenient return for prefixes which do not appear in map. */
 261                 static const Map empty_map;
 262                 return &empty_map;
 263         }
 264
 265         return top;
 266 }
 267
 268 static void clear(Map n)
 269 {
 270         if (!n.v) {
 271                 clear(n.u.n->child[0]);
 272                 clear(n.u.n->child[1]);
 273                 free(n.u.n);
 274         } else {
 275                 free((char*)n.u.s);
 276         }
 277 }
 278
 279 void map_clear(Map *map)
 280 {
 281         if (map->u.n)
 282                 clear(*map);
 283         map->u.n = NULL;
 284         map->v = NULL;
 285 }
 286
 287 static bool copy(Map *dest, Map n)
 288 {
 289         if (!n.v) {
 290                 return copy(dest, n.u.n->child[0]) &&
 291                        copy(dest, n.u.n->child[1]);
 292         } else {
 293                 if (!map_put(dest, n.u.s, n.v) && map_get(dest, n.u.s) != n.v) {
 294                         map_delete(dest, n.u.s);
 295                         return map_put(dest, n.u.s, n.v);
 296                 }
 297                 return true;
 298         }
 299 }
 300
 301 bool map_copy(Map *dest, Map *src)
 302 {
 303         if (!src || !src->u.n)
 304                 return true;
 305
 306         return copy(dest, *src);
 307 }
 308
 309 bool map_empty(const Map *map)
 310 {
 311         return map->u.n == NULL;
 312 }
 313
 314 Map *map_new(void)
 315 {
 316         return calloc(1, sizeof(Map));
 317 }
 318
 319 void map_free(Map *map)
 320 {
 321         if (!map)
 322                 return;
 323         map_clear(map);
 324         free(map);
 325 }