6bf4cf9c6768ca03cf5f32d638122e2742fe5e02
[apps/pfixtools.git] / trie.c
1 /******************************************************************************/
2 /*          pfixtools: a collection of postfix related tools                  */
3 /*          ~~~~~~~~~                                                         */
4 /*  ________________________________________________________________________  */
5 /*                                                                            */
6 /*  Redistribution and use in source and binary forms, with or without        */
7 /*  modification, are permitted provided that the following conditions        */
8 /*  are met:                                                                  */
9 /*                                                                            */
10 /*  1. Redistributions of source code must retain the above copyright         */
11 /*     notice, this list of conditions and the following disclaimer.          */
12 /*  2. Redistributions in binary form must reproduce the above copyright      */
13 /*     notice, this list of conditions and the following disclaimer in the    */
14 /*     documentation and/or other materials provided with the distribution.   */
15 /*  3. The names of its contributors may not be used to endorse or promote    */
16 /*     products derived from this software without specific prior written     */
17 /*     permission.                                                            */
18 /*                                                                            */
19 /*  THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND   */
20 /*  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE     */
21 /*  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR        */
22 /*  PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS    */
23 /*  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR    */
24 /*  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF      */
25 /*  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS  */
26 /*  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN   */
27 /*  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)   */
28 /*  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF    */
29 /*  THE POSSIBILITY OF SUCH DAMAGE.                                           */
30 /******************************************************************************/
31
32 /*
33  * Copyright © 2008 Florent Bruneau
34  */
35
36 #include "array.h"
37 #include "str.h"
38 #include "trie.h"
39
40 typedef struct trie_entry_t trie_entry_t;
41
42 struct trie_entry_t {
43     uint32_t c_offset;
44     uint32_t children_offset;
45
46     uint16_t c_len;
47     uint16_t children_len;
48 };
49 #define TRIE_ENTRY_INIT { 0, 0, 0, 0 }
50 ARRAY(trie_entry_t)
51
52 #define str(trie, entry)  array_ptr((trie)->c, (entry)->c_offset)
53 #define key(trie, id) array_ptr((trie)->keys, array_elt((trie)->keys_offset, (id)))
54
55 struct trie_t {
56     A(trie_entry_t) entries;
57     A(char)         c;
58     A(char)         keys;
59     A(int)          keys_offset;
60
61     bool locked;
62 };
63
64 trie_t *trie_new(void)
65 {
66     return p_new(trie_t, 1);
67 }
68
69 static inline void trie_cleanup_build_data(trie_t *trie)
70 {
71     array_wipe(trie->keys);
72     array_wipe(trie->keys_offset);
73 }
74
75 void trie_delete(trie_t **trie)
76 {
77     if (*trie) {
78         trie_cleanup_build_data(*trie);
79         trie_unlock(*trie);
80         array_wipe((*trie)->entries);
81         array_wipe((*trie)->c);
82         p_delete(trie);
83     }
84 }
85
86 /** Check that the given entry is a prefix for the given key.
87  */
88 static inline bool trie_entry_c_match(const trie_t *trie,
89                                       const trie_entry_t *entry,
90                                       const char *key)
91 {
92     const char *c = str(trie, entry);
93     int i = 0;
94     for (i = 0 ; i < entry->c_len ; ++i) {
95         if (key[i] != c[i]) {
96             return false;
97         }
98     }
99     return true;
100 }
101
102 static inline bool trie_entry_match(const trie_t *trie,
103                                     const trie_entry_t *entry, const char *key)
104 {
105     return !!(strcmp(str(trie, entry), key) == 0);
106 }
107
108 static inline bool trie_entry_prefix(const trie_t *trie,
109                                      const trie_entry_t *entry, const char *key)
110 {
111     int len = entry->c_len;
112     if (len > 0 && str(trie, entry)[len -1] == '\0') {
113         --len;
114     }
115     return !!(strncmp(str(trie, entry), key, len) == 0);
116 }
117
118 static inline bool trie_entry_is_leaf(const trie_entry_t *entry)
119 {
120     return entry->children_len == 0;
121 }
122
123 /** Lookup for a child of entry matching the given entry at the given pos.
124  * Only the first character of the children is taken into account in the
125  * lookup. The current entry is assumed to match the key.
126  */
127 static inline const trie_entry_t* trie_entry_child(const trie_t *trie,
128                                                    const trie_entry_t* entry,
129                                                    const char *key)
130 {
131     uint32_t start = entry->children_offset;
132     uint32_t end   = start + entry->children_len;
133     const char c = *key;
134
135     while (start < end) {
136         uint32_t mid = (start + end) >> 1;
137         const trie_entry_t* child = array_ptr(trie->entries, mid);
138         const char c2 = str(trie, child)[0];
139
140         if (c2 == c) {
141           return child;
142         }
143         if (c < c2) {
144           end = mid;
145         } else {
146           start = mid + 1;
147         }
148     }
149     return NULL;
150 }
151
152 static inline uint32_t trie_entry_new(trie_t *trie)
153 {
154     const trie_entry_t e = TRIE_ENTRY_INIT;
155     array_add(trie->entries, e);
156     return trie->entries.len - 1;
157 }
158
159 static inline uint32_t trie_add_leaf(trie_t *trie, const char *key)
160 {
161     trie_entry_t *entry;
162     int len = m_strlen(key) + 1;
163     int id  = trie_entry_new(trie);
164     entry = array_ptr(trie->entries, id);
165     entry->c_offset = trie->c.len;
166     entry->c_len    = len;
167 #ifdef CHECK_INTEGRITY
168     for (int i = 0 ; i < len - 1 ; ++i) {
169         if (key[i] == '\0') {
170             printf("Found a '\\0' in the string of the leaf\n");
171             abort();
172         }
173     }
174     if (key[len - 1] != '\0') {
175       printf("Key does not end with a '\\0'");
176       abort();
177     }
178 #endif
179     array_append(trie->c, key, len);
180     return trie->entries.len - 1;
181 }
182
183 static inline void trie_entry_insert_child(trie_t *trie, uint32_t id, uint32_t pchild)
184 {
185     trie_entry_t *entry = array_ptr(trie->entries, id);
186     if (entry->children_len == 0) {
187         entry->children_offset = pchild;
188         entry->children_len    = 1;
189     } else {
190         if (entry->children_offset + entry->children_len != pchild) {
191             printf("Inserting child %d while offset is %d[%d]\n",
192                    pchild, entry->children_offset, entry->children_len);
193             abort();
194         }
195         ++entry->children_len;
196     }
197 }
198
199 static inline void trie_entry_split(trie_t *trie, uint32_t id, uint16_t pos)
200 {
201     trie_entry_t *child;
202     trie_entry_t *entry;
203     child    = array_ptr(trie->entries, trie_entry_new(trie));
204     entry    = array_ptr(trie->entries, id);
205     if (pos == 0) {
206         child->c_offset = entry->c_offset;
207         child->c_len    = entry->c_len;
208         entry->c_offset = 0;
209         entry->c_len    = 0;
210     } else {
211         assert(pos <= entry->c_len);
212         child->c_offset = entry->c_offset + pos;
213         child->c_len    = entry->c_len - pos;
214         entry->c_len    = pos;
215     }
216     child->children_offset = entry->children_offset;
217     child->children_len    = entry->children_len;
218     entry->children_offset = trie->entries.len - 1;
219     entry->children_len    = 1;
220 }
221
222 void trie_insert(trie_t *trie, const char* key)
223 {
224     assert(trie->entries.len == 0 && "Trie already compiled");
225
226     int len = m_strlen(key) + 1;
227     array_add(trie->keys_offset, trie->keys.len);
228     array_append(trie->keys, key, len);
229 }
230
231
232 static inline void trie_compile_aux(trie_t *trie, uint32_t id,
233                                     uint32_t first_key, uint32_t last_key,
234                                     int offset, int initial_diff)
235 {
236     uint32_t forks[256];
237     uint32_t fork_pos = 0;
238     char current = '\0';
239
240 #ifdef CHECK_INTEGRITY
241     assert(strcmp(key(trie, first_key) + offset, str(trie, entry)) == 0);
242 #endif
243
244     for (int off_diff = initial_diff ; fork_pos == 0 ; ++off_diff, ++offset) {
245         current = key(trie, first_key)[offset];
246         for (uint32_t i = first_key + 1 ; i < last_key ; ++i) {
247             const char *ckey = key(trie, i) + offset;
248             const char c = *ckey;
249             if (c != current) {
250                 array_ensure_capacity_delta(trie->entries, 2);
251                 if (fork_pos == 0) {
252                     trie_entry_split(trie, id, off_diff);
253                 }
254                 trie_entry_insert_child(trie, id, trie_add_leaf(trie, ckey));
255                 forks[fork_pos++] = i;
256                 current = c;
257             }
258         }
259         if (fork_pos == 0 && current == '\0') {
260             return;
261         }
262     }
263     forks[fork_pos] = last_key;
264
265     const uint8_t children_len = array_elt(trie->entries, id).children_len;
266     for (uint16_t i = 0 ; i < children_len ; ++i) {
267         int child = array_elt(trie->entries, id).children_offset + i;
268         if (forks[i] - 1 > first_key) {
269             trie_compile_aux(trie, child, first_key, forks[i], offset, 1);
270         }
271         first_key = forks[i];
272     }
273 }
274
275 void trie_compile(trie_t *trie, bool memlock)
276 {
277     assert(trie->entries.len == 0 && "Trie already compiled");
278     assert(trie->keys.len != 0 && "Trying to compile an empty trie");
279     {
280 #       define QSORT_TYPE int
281 #       define QSORT_BASE trie->keys_offset.data
282 #       define QSORT_NELT trie->keys_offset.len
283 #       define QSORT_LT(a,b) strcmp(trie->keys.data + *a, trie->keys.data + *b) < 0
284 #       include "qsort.c"
285     }
286
287     array_ensure_capacity(trie->entries, trie->keys_offset.len);
288     trie_compile_aux(trie, trie_add_leaf(trie, key(trie, 0)),
289                      0, trie->keys_offset.len, 0, 0);
290     trie_cleanup_build_data(trie);
291     array_adjust(trie->entries);
292     array_adjust(trie->c);
293     if (memlock) {
294         trie_lock(trie);
295     }
296 }
297
298 bool trie_lookup(const trie_t *trie, const char *key)
299 {
300     assert(trie->keys.len == 0L && "Can't lookup: trie not compiled");
301     if (trie->entries.len == 0) {
302         return false;
303     } else {
304         const trie_entry_t *current = array_ptr(trie->entries, 0);
305         while (true) {
306             if (trie_entry_is_leaf(current)) {
307                 return trie_entry_match(trie, current, key);
308             } else if (trie_entry_c_match(trie, current, key)) {
309                 key += current->c_len;
310                 current = trie_entry_child(trie, current, key);
311                 if (current == NULL) {
312                     return false;
313                 }
314             } else {
315                 return false;
316             }
317         }
318     }
319 }
320
321 bool trie_prefix(const trie_t *trie, const char *key)
322 {
323     assert(trie->keys.len == 0L && "Can't lookup: trie not compiled");
324     if (trie->entries.len == 0) {
325         return false;
326     } else {
327         const trie_entry_t *current = array_ptr(trie->entries, 0);
328         while (true) {
329             if (trie_entry_is_leaf(current)) {
330                 return trie_entry_prefix(trie, current, key);
331             } else if (trie_entry_c_match(trie, current, key)) {
332                 key += current->c_len;
333                 current = trie_entry_child(trie, current, key);
334                 if (current == NULL) {
335                     return false;
336                 }
337             } else {
338                 return false;
339             }
340         }
341     }
342 }
343
344 void trie_lock(trie_t *trie)
345 {
346     if (trie->locked) {
347         return;
348     }
349     if (!array_lock(trie->entries)) {
350         UNIXERR("mlock");
351     }
352     if (!array_lock(trie->c)) {
353         UNIXERR("mlock");
354     }
355     if (mlock(trie, sizeof(trie_t)) != 0) {
356         UNIXERR("mlock");
357         return;
358     }
359     trie->locked = true;
360 }
361
362 void trie_unlock(trie_t *trie)
363 {
364     if (!trie->locked) {
365         return;
366     }
367     array_unlock(trie->entries);
368     array_unlock(trie->c);
369     munlock(trie, sizeof(trie_t));
370     trie->locked = false;
371 }
372
373 /* Debug {{{1
374  */
375
376 static inline void trie_entry_inspect(const trie_t *trie, bool show_content,
377                                       const trie_entry_t *entry, int level)
378 {
379     static int max_depth = 0;
380     static int leaves    = 0;
381     static int depth_sum = 0;
382
383     if (entry == array_ptr(trie->entries, 0)) {
384       max_depth = 0;
385       leaves    = 0;
386       depth_sum = 0;
387     }
388     if (trie_entry_is_leaf(entry)) {
389         if (level > max_depth) {
390             max_depth = level;
391         }
392         ++leaves;
393         depth_sum += level;
394     }
395     if (show_content) {
396         for (int i = 0 ; i < level ; ++i) {
397             fputs("  ", stdout);
398         }
399         if (entry->c_len == 0) {
400             fputs("(0)", stdout);
401         } else {
402             const char *c = str(trie, entry);
403             printf("(%d) ", entry->c_len);
404             for (int i = 0 ; i < entry->c_len ; ++i) {
405                 if (c[i]) {
406                     printf("%c ", c[i]);
407                 } else {
408                     fputs("\\0 ", stdout);
409                 }
410             }
411         }
412         fputs("\n", stdout);
413     }
414     for (uint32_t i = 0 ; i < entry->children_len ; ++i) {
415         trie_entry_inspect(trie, show_content,
416                            array_ptr(trie->entries, entry->children_offset + i),
417                            level + 1);
418     }
419     if (level == 0) {
420         printf("Average char per node: %d\n", trie->c.len / trie->entries.len);
421         printf("Number of nodes: %d\n", trie->entries.len);
422         printf("Number of leaves: %d\n", leaves);
423         printf("Max depth: %d\n", max_depth);
424         printf("Average leaf depth: %d\n", depth_sum / leaves);
425         printf("Memory used: %zd\n", (trie->entries.size * sizeof(trie_entry_t))
426                                   + (trie->c.size) + sizeof(trie_t));
427     }
428 }
429
430 void trie_inspect(const trie_t *trie, bool show_content)
431 {
432     trie_entry_inspect(trie, show_content, trie->entries.data, 0);
433 }