git.ozlabs.org Git - ccan/blob - ccan/strset/strset.c

   1 /* This code is based on the public domain code at
   2  * http://github.com/agl/critbit writtem by Adam Langley
   3  * <agl@imperialviolet.org>.
   4  *
   5  * Here are the main implementation differences:
   6  * (1) We don't strdup the string on insert; we use the pointer we're given.
   7  * (2) We use a straight bit number rather than a mask; it's simpler.
   8  * (3) We don't use the bottom bit of the pointer, but instead use a leading
   9  *     zero to distinguish nodes from strings.
  10  * (4) The empty string (which would look like a node) is handled
  11  *     using a special "empty node".
  12  * (5) Delete returns the string, so you can free it if you want to.
  13  * (6) Unions instead of void *, bool instead of int.
  14  */
  15 #include <ccan/strset/strset.h>
  16 #include <ccan/short_types/short_types.h>
  17 #include <ccan/likely/likely.h>
  18 #include <ccan/str/str.h>
  19 #include <ccan/ilog/ilog.h>
  20 #include <assert.h>
  21 #include <stdlib.h>
  22 #include <errno.h>
  23
  24 struct node {
  25         /* To differentiate us from strings. */
  26         char nul_byte;
  27         /* The bit where these children differ. */
  28         u8 bit_num;
  29         /* The byte number where first bit differs (-1 == empty string node). */
  30         size_t byte_num;
  31         /* These point to strings or nodes. */
  32         struct strset child[2];
  33 };
  34
  35 /* Closest member to this in a non-empty set. */
  36 static const char *closest(struct strset n, const char *member)
  37 {
  38         size_t len = strlen(member);
  39         const u8 *bytes = (const u8 *)member;
  40
  41         /* Anything with first byte 0 is a node. */
  42         while (!n.u.s[0]) {
  43                 u8 direction = 0;
  44
  45                 /* Special node which represents the empty string. */
  46                 if (unlikely(n.u.n->byte_num == (size_t)-1)) {
  47                         n = n.u.n->child[0];
  48                         break;
  49                 }
  50
  51                 if (n.u.n->byte_num < len) {
  52                         u8 c = bytes[n.u.n->byte_num];
  53                         direction = (c >> n.u.n->bit_num) & 1;
  54                 }
  55                 n = n.u.n->child[direction];
  56         }
  57         return n.u.s;
  58 }
  59
  60 char *strset_get(const struct strset *set, const char *member)
  61 {
  62         const char *str;
  63
  64         /* Non-empty set? */
  65         if (set->u.n) {
  66                 str = closest(*set, member);
  67                 if (streq(member, str))
  68                         return (char *)str;
  69         }
  70         errno = ENOENT;
  71         return NULL;
  72 }
  73
  74 static bool set_string(struct strset *set,
  75                        struct strset *n, const char *member)
  76 {
  77         /* Substitute magic empty node if this is the empty string */
  78         if (unlikely(!member[0])) {
  79                 n->u.n = malloc(sizeof(*n->u.n));
  80                 if (unlikely(!n->u.n)) {
  81                         errno = ENOMEM;
  82                         return false;
  83                 }
  84                 n->u.n->nul_byte = '\0';
  85                 n->u.n->byte_num = (size_t)-1;
  86                 /* Attach the string to child[0] */
  87                 n = &n->u.n->child[0];
  88         }
  89         n->u.s = member;
  90         return true;
  91 }
  92
  93 bool strset_add(struct strset *set, const char *member)
  94 {
  95         size_t len = strlen(member);
  96         const u8 *bytes = (const u8 *)member;
  97         struct strset *np;
  98         const char *str;
  99         struct node *newn;
 100         size_t byte_num;
 101         u8 bit_num, new_dir;
 102
 103         /* Empty set? */
 104         if (!set->u.n) {
 105                 return set_string(set, set, member);
 106         }
 107
 108         /* Find closest existing member. */
 109         str = closest(*set, member);
 110
 111         /* Find where they differ. */
 112         for (byte_num = 0; str[byte_num] == member[byte_num]; byte_num++) {
 113                 if (member[byte_num] == '\0') {
 114                         /* All identical! */
 115                         errno = EEXIST;
 116                         return false;
 117                 }
 118         }
 119
 120         /* Find which bit differs (if we had ilog8, we'd use it) */
 121         bit_num = ilog32_nz((u8)str[byte_num] ^ bytes[byte_num]) - 1;
 122         assert(bit_num < CHAR_BIT);
 123
 124         /* Which direction do we go at this bit? */
 125         new_dir = ((bytes[byte_num]) >> bit_num) & 1;
 126
 127         /* Allocate new node. */
 128         newn = malloc(sizeof(*newn));
 129         if (!newn) {
 130                 errno = ENOMEM;
 131                 return false;
 132         }
 133         newn->nul_byte = '\0';
 134         newn->byte_num = byte_num;
 135         newn->bit_num = bit_num;
 136         if (unlikely(!set_string(set, &newn->child[new_dir], member))) {
 137                 free(newn);
 138                 return false;
 139         }
 140
 141         /* Find where to insert: not closest, but first which differs! */
 142         np = set;
 143         while (!np->u.s[0]) {
 144                 u8 direction = 0;
 145
 146                 /* Special node which represents the empty string will
 147                  * break here too! */
 148                 if (np->u.n->byte_num > byte_num)
 149                         break;
 150                 /* Subtle: bit numbers are "backwards" for comparison */
 151                 if (np->u.n->byte_num == byte_num && np->u.n->bit_num < bit_num)
 152                         break;
 153
 154                 if (np->u.n->byte_num < len) {
 155                         u8 c = bytes[np->u.n->byte_num];
 156                         direction = (c >> np->u.n->bit_num) & 1;
 157                 }
 158                 np = &np->u.n->child[direction];
 159         }
 160
 161         newn->child[!new_dir]= *np;
 162         np->u.n = newn;
 163         return true;
 164 }
 165
 166 char *strset_del(struct strset *set, const char *member)
 167 {
 168         size_t len = strlen(member);
 169         const u8 *bytes = (const u8 *)member;
 170         struct strset *parent = NULL, *n;
 171         const char *ret = NULL;
 172         u8 direction = 0; /* prevent bogus gcc warning. */
 173
 174         /* Empty set? */
 175         if (!set->u.n) {
 176                 errno = ENOENT;
 177                 return NULL;
 178         }
 179
 180         /* Find closest, but keep track of parent. */
 181         n = set;
 182         /* Anything with first byte 0 is a node. */
 183         while (!n->u.s[0]) {
 184                 u8 c = 0;
 185
 186                 /* Special node which represents the empty string. */
 187                 if (unlikely(n->u.n->byte_num == (size_t)-1)) {
 188                         const char *empty_str = n->u.n->child[0].u.s;
 189
 190                         if (member[0]) {
 191                                 errno = ENOENT;
 192                                 return NULL;
 193                         }
 194
 195                         /* Sew empty string back so remaining logic works */
 196                         free(n->u.n);
 197                         n->u.s = empty_str;
 198                         break;
 199                 }
 200
 201                 parent = n;
 202                 if (n->u.n->byte_num < len) {
 203                         c = bytes[n->u.n->byte_num];
 204                         direction = (c >> n->u.n->bit_num) & 1;
 205                 } else
 206                         direction = 0;
 207                 n = &n->u.n->child[direction];
 208         }
 209
 210         /* Did we find it? */
 211         if (!streq(member, n->u.s)) {
 212                 errno = ENOENT;
 213                 return NULL;
 214         }
 215
 216         ret = n->u.s;
 217
 218         if (!parent) {
 219                 /* We deleted last node. */
 220                 set->u.n = NULL;
 221         } else {
 222                 struct node *old = parent->u.n;
 223                 /* Raise other node to parent. */
 224                 *parent = old->child[!direction];
 225                 free(old);
 226         }
 227
 228         return (char *)ret;
 229 }
 230
 231 static bool iterate(struct strset n,
 232                     bool (*handle)(const char *, void *), const void *data)
 233 {
 234         if (n.u.s[0])
 235                 return handle(n.u.s, (void *)data);
 236         if (unlikely(n.u.n->byte_num == (size_t)-1))
 237                 return handle(n.u.n->child[0].u.s, (void *)data);
 238
 239         return iterate(n.u.n->child[0], handle, data)
 240                 && iterate(n.u.n->child[1], handle, data);
 241 }
 242
 243 void strset_iterate_(const struct strset *set,
 244                      bool (*handle)(const char *, void *), const void *data)
 245 {
 246         /* Empty set? */
 247         if (!set->u.n)
 248                 return;
 249
 250         iterate(*set, handle, data);
 251 }
 252
 253 const struct strset *strset_prefix(const struct strset *set, const char *prefix)
 254 {
 255         const struct strset *n, *top;
 256         size_t len = strlen(prefix);
 257         const u8 *bytes = (const u8 *)prefix;
 258
 259         /* Empty set -> return empty set. */
 260         if (!set->u.n)
 261                 return set;
 262
 263         top = n = set;
 264
 265         /* We walk to find the top, but keep going to check prefix matches. */
 266         while (!n->u.s[0]) {
 267                 u8 c = 0, direction;
 268
 269                 /* Special node which represents the empty string. */
 270                 if (unlikely(n->u.n->byte_num == (size_t)-1)) {
 271                         n = &n->u.n->child[0];
 272                         break;
 273                 }
 274
 275                 if (n->u.n->byte_num < len)
 276                         c = bytes[n->u.n->byte_num];
 277
 278                 direction = (c >> n->u.n->bit_num) & 1;
 279                 n = &n->u.n->child[direction];
 280                 if (c)
 281                         top = n;
 282         }
 283
 284         if (!strstarts(n->u.s, prefix)) {
 285                 /* Convenient return for prefixes which do not appear in set. */
 286                 static const struct strset empty_set;
 287                 return &empty_set;
 288         }
 289
 290         return top;
 291 }
 292
 293 static void clear(struct strset n)
 294 {
 295         if (!n.u.s[0]) {
 296                 if (likely(n.u.n->byte_num != (size_t)-1)) {
 297                         clear(n.u.n->child[0]);
 298                         clear(n.u.n->child[1]);
 299                 }
 300                 free(n.u.n);
 301         }
 302 }
 303
 304 void strset_clear(struct strset *set)
 305 {
 306         if (set->u.n)
 307                 clear(*set);
 308         set->u.n = NULL;
 309 }