htable: add allocator hooks.

Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>

diff --git a/ccan/htable/htable.c b/ccan/htable/htable.c
index f3568230af6136e99da51d1434a435d51ac843cf..e6dc4fd817f886fa2676d0cefd9774f0fd9a690a 100644 (file)
--- a/ccan/htable/htable.c
+++ b/ccan/htable/htable.c
@@ -11,6 +11,30 @@
 /* We use 0x1 as deleted marker. */
 #define HTABLE_DELETED (0x1)
 
 /* We use 0x1 as deleted marker. */
 #define HTABLE_DELETED (0x1)
 

+static void *htable_default_alloc(struct htable *ht, size_t len)
+{
+       return calloc(len, 1);
+}
+
+static void htable_default_free(struct htable *ht, void *p)
+{
+       free(p);
+}
+
+static void *(*htable_alloc)(struct htable *, size_t) = htable_default_alloc;
+static void (*htable_free)(struct htable *, void *) = htable_default_free;
+
+void htable_set_allocator(void *(*alloc)(struct htable *, size_t len),
+                         void (*free)(struct htable *, void *p))
+{
+       if (!alloc)
+               alloc = htable_default_alloc;
+       if (!free)
+               free = htable_default_free;
+       htable_alloc = alloc;
+       htable_free = free;
+}
+

 /* We clear out the bits which are always the same, and put metadata there. */
 static inline uintptr_t get_extra_ptr_bits(const struct htable *ht,
                                           uintptr_t e)
 /* We clear out the bits which are always the same, and put metadata there. */
 static inline uintptr_t get_extra_ptr_bits(const struct htable *ht,
                                           uintptr_t e)


@@ -73,7 +97,7 @@ bool htable_init_sized(struct htable *ht,
                        break;
        }
 
                        break;
        }
 

-       ht->table = calloc(1 << ht->bits, sizeof(size_t));
+       ht->table = htable_alloc(ht, sizeof(size_t) << ht->bits);

        if (!ht->table) {
                ht->table = &ht->perfect_bit;
                return false;
        if (!ht->table) {
                ht->table = &ht->perfect_bit;
                return false;


@@ -86,13 +110,13 @@ bool htable_init_sized(struct htable *ht,
 void htable_clear(struct htable *ht)
 {
        if (ht->table != &ht->perfect_bit)
 void htable_clear(struct htable *ht)
 {
        if (ht->table != &ht->perfect_bit)

-               free((void *)ht->table);
+               htable_free(ht, (void *)ht->table);

        htable_init(ht, ht->rehash, ht->priv);
 }
 
 bool htable_copy_(struct htable *dst, const struct htable *src)
 {
        htable_init(ht, ht->rehash, ht->priv);
 }
 
 bool htable_copy_(struct htable *dst, const struct htable *src)
 {

-       uintptr_t *htable = malloc(sizeof(size_t) << src->bits);
+       uintptr_t *htable = htable_alloc(dst, sizeof(size_t) << src->bits);

 
        if (!htable)
                return false;
 
        if (!htable)
                return false;


@@ -189,7 +213,7 @@ static COLD bool double_table(struct htable *ht)
        uintptr_t *oldtable, e;
 
        oldtable = ht->table;
        uintptr_t *oldtable, e;
 
        oldtable = ht->table;

-       ht->table = calloc(1 << (ht->bits+1), sizeof(size_t));
+       ht->table = htable_alloc(ht, sizeof(size_t) << (ht->bits+1));

        if (!ht->table) {
                ht->table = oldtable;
                return false;
        if (!ht->table) {
                ht->table = oldtable;
                return false;


@@ -214,7 +238,7 @@ static COLD bool double_table(struct htable *ht)
                                ht_add(ht, p, ht->rehash(p, ht->priv));
                        }
                }
                                ht_add(ht, p, ht->rehash(p, ht->priv));
                        }
                }

-               free(oldtable);
+               htable_free(ht, oldtable);

        }
        ht->deleted = 0;
 
        }
        ht->deleted = 0;
 

diff --git a/ccan/htable/htable.h b/ccan/htable/htable.h
index 53c447c09b952b198913efa25a48b82d34777242..28755d617d67b7766d44fe3e5ce93d7956a85d55 100644 (file)
--- a/ccan/htable/htable.h
+++ b/ccan/htable/htable.h
@@ -259,4 +259,11 @@ void *htable_prev_(const struct htable *htable, struct htable_iter *i);
        htable_delval_(htable_debug(htable, HTABLE_LOC), i)
 void htable_delval_(struct htable *ht, struct htable_iter *i);
 
        htable_delval_(htable_debug(htable, HTABLE_LOC), i)
 void htable_delval_(struct htable *ht, struct htable_iter *i);
 

+/**
+ * htable_set_allocator - set calloc/free functions.
+ * @alloc: allocator to use, must zero memory!
+ * @free: unallocator to use (@p is NULL or a return from @alloc)
+ */
+void htable_set_allocator(void *(*alloc)(struct htable *, size_t len),
+                         void (*free)(struct htable *, void *p));

 #endif /* CCAN_HTABLE_H */
 #endif /* CCAN_HTABLE_H */

diff --git a/ccan/htable/test/run-allocator.c b/ccan/htable/test/run-allocator.c
new file mode 100644 (file)
index 0000000..ffbb48d
--- /dev/null
+++ b/ccan/htable/test/run-allocator.c
@@ -0,0 +1,70 @@
+/* Include the C files directly. */
+#include <ccan/htable/htable.h>
+#include <ccan/htable/htable.c>
+#include <ccan/tap/tap.h>
+#include <stdbool.h>
+#include <string.h>
+
+struct htable_with_counters {
+       struct htable ht;
+       size_t num_alloc, num_free;
+};
+
+static void *test_alloc(struct htable *ht, size_t len)
+{
+       ((struct htable_with_counters *)ht)->num_alloc++;
+       return calloc(len, 1);
+}
+
+       
+static void test_free(struct htable *ht, void *p)
+{
+       if (p) {
+               ((struct htable_with_counters *)ht)->num_free++;
+               free(p);
+       }
+}
+
+static size_t hash(const void *elem, void *unused UNNEEDED)
+{
+       return *(size_t *)elem;
+}
+
+int main(void)
+{
+       struct htable_with_counters htc;
+       size_t val[] = { 0, 1 };
+
+       htc.num_alloc = htc.num_free = 0;
+       plan_tests(12);
+
+       htable_set_allocator(test_alloc, test_free);
+       htable_init(&htc.ht, hash, NULL);
+       htable_add(&htc.ht, hash(&val[0], NULL), &val[0]);
+       ok1(htc.num_alloc == 1);
+       ok1(htc.num_free == 0);
+       /* Adding another increments, then frees old */
+       htable_add(&htc.ht, hash(&val[1], NULL), &val[1]);
+       ok1(htc.num_alloc == 2);
+       ok1(htc.num_free == 1);
+       htable_clear(&htc.ht);
+       ok1(htc.num_alloc == 2);
+       ok1(htc.num_free == 2);
+
+       /* Should restore defaults */
+       htable_set_allocator(NULL, NULL);
+       ok1(htable_alloc == htable_default_alloc);
+       ok1(htable_free == htable_default_free);
+
+       htable_init(&htc.ht, hash, NULL);
+       htable_add(&htc.ht, hash(&val[0], NULL), &val[0]);
+       ok1(htc.num_alloc == 2);
+       ok1(htc.num_free == 2);
+       htable_add(&htc.ht, hash(&val[1], NULL), &val[1]);
+       ok1(htc.num_alloc == 2);
+       ok1(htc.num_free == 2);
+       htable_clear(&htc.ht);
+
+       /* This exits depending on whether all tests passed */
+       return exit_status();
+}