1 files changed, 21 insertions, 15 deletions

diff --git a/util/dict.c b/util/dict.c
index 75fcf3b..d50b07b 100644
--- a/util/dict.c
+++ b/util/dict.c
@@ -4,7 +4,7 @@
 
 #include "util/dict.h"
 
-#define CHAR_TO_PTR(c) (bit_to_ptr[char_to_bit[c]])
+#define CHAR_TO_PTR(c) (bit_to_ptr[char_to_bit(c)])
 #define popcount(x) (__builtin_popcount(x))
 
 #define start_level d->start_level
@@ -12,17 +12,19 @@
 #define num_levels d->num_levels
 #define strings d->strings
 #define nstrings d->nstrings
-#define char_to_bit d->char_to_bit
+#define char_to_bit(c) d->char_to_bit[(size_t)(c)]
 
 int dict_compile(struct dict *d)
 {
+    void *(*new_calloc)(size_t, size_t) = d->calloc ? d->calloc : calloc;
+
     // max number of levels
     for(size_t i = 0; i < nstrings; i++)
         if(strlen(strings[i].s) > num_levels) num_levels = strlen(strings[i].s);
 
     // allocated levels
     for(size_t i = 0; i < MAPPED_CHARS; i++) {
-        bit_to_ptr[i] = calloc(num_levels, sizeof(*bit_to_ptr[i]));
+        bit_to_ptr[i] = new_calloc(num_levels, sizeof(*bit_to_ptr[i]));
         if(!bit_to_ptr[i]) return 1;
     }
 
@@ -32,7 +34,7 @@ int dict_compile(struct dict *d)
     for(size_t i = 0; i < nstrings; i++) {
         struct level *l = &start_level;
         for(size_t j = 0; j < strlen(strings[i].s)+1; j++) {
-            uint8_t bit = char_to_bit[strings[i].s[j]];
+            uint8_t bit = char_to_bit(strings[i].s[j]);
 
             l->bit_mask |= 1 << bit;
             l = &bit_to_ptr[bit][j];
@@ -44,10 +46,12 @@ int dict_compile(struct dict *d)
     for(size_t i = 0; i < MAPPED_CHARS; i++) {
         struct level *l = &start_level;
         for(size_t j = 0; j < num_levels + 1; j++) {
-            l->token_masks = realloc(l->token_masks, popcount(l->bit_mask)
-                                    * sizeof(*l->token_masks));
-            memset(l->token_masks, 0, popcount(l->bit_mask)
-                   * sizeof(*l->token_masks));
+            if(!l->token_masks) {
+                l->token_masks = new_calloc(popcount(l->bit_mask), sizeof(*l->token_masks));
+                if(!l->token_masks) return 1;
+                memset(l->token_masks, 0, popcount(l->bit_mask)
+                       * sizeof(*l->token_masks));
+            }
 
             l = &bit_to_ptr[i][j];
         }
@@ -57,7 +61,7 @@ int dict_compile(struct dict *d)
     for(size_t i = 0; i < nstrings; i++) {
         struct level *l = &start_level;
         for(size_t j = 0; j < strlen(strings[i].s)+1; j++) {
-            uint8_t bit = char_to_bit[strings[i].s[j]];
+            uint8_t bit = char_to_bit(strings[i].s[j]);
             uint8_t idx = popcount(l->bit_mask & ((1 << bit) - 1));
 
             l->token_masks[idx] |= 1 << strings[i].t;
@@ -93,13 +97,15 @@ void dict_print(struct dict *d)
 
 void dict_free(struct dict *d)
 {
-    free(start_level.token_masks);
+    void (*new_free)(void *) = d->free ? d->free : free;
+
+    new_free(start_level.token_masks);
     for(size_t i = 0; i < MAPPED_CHARS; i++) {
         for(size_t j = 0; j < num_levels; j++) {
             if(bit_to_ptr[i][j].token_masks)
-                free(bit_to_ptr[i][j].token_masks);
+                new_free(bit_to_ptr[i][j].token_masks);
         }
-        free(bit_to_ptr[i]);
+        new_free(bit_to_ptr[i]);
     }
 }
 
@@ -110,9 +116,9 @@ int dict_check(struct dict *d, char *string)
     for(size_t i = 0; i < strlen(string) + 1; i++) {
         struct level *l = (i == 0)
             ? &start_level
-            : &bit_to_ptr[char_to_bit[string[i-1]]][i-1];
+            : &bit_to_ptr[char_to_bit(string[i-1])][i-1];
 
-        uint8_t bit = char_to_bit[string[i]];
+        uint8_t bit = char_to_bit(string[i]);
 
         if((l->bit_mask & (1 << bit)) == 0) return -1;
 
@@ -175,7 +181,7 @@ int main(void)
     d.strings = strings;
     d.nstrings = nstrings;
     d.char_to_bit = char_to_bit;
-        
+
     dict_compile(&d);
 
     int t;