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