Add custom allocator support

18 files changed, 438 insertions, 110 deletions

diff --git a/benchmark/dict_bench.c b/benchmark/dict_bench.c
index e1583f3..cd358a3 100644
--- a/benchmark/dict_bench.c
+++ b/benchmark/dict_bench.c
@@ -139,8 +139,10 @@ main(int argc, char** argv)
     printf("Benchmarking n = %zu\n", n);
     GHashTable* hash = g_hash_table_new(g_str_hash, g_str_equal);
 
-    ZixHash* zhash = zix_hash_new(
-      identity, (ZixHashFunc)zix_chunk_hash, (ZixEqualFunc)zix_chunk_equal);
+    ZixHash* zhash = zix_hash_new(NULL,
+                                  identity,
+                                  (ZixHashFunc)zix_chunk_hash,
+                                  (ZixEqualFunc)zix_chunk_equal);
 
     fprintf(insert_dat, "%zu", n);
     fprintf(search_dat, "%zu", n);
diff --git a/benchmark/tree_bench.c b/benchmark/tree_bench.c
index 6e573ee..b55d772 100644
--- a/benchmark/tree_bench.c
+++ b/benchmark/tree_bench.c
@@ -90,7 +90,7 @@ bench_zix_tree(size_t n_elems,
 
   uintptr_t    r  = 0;
   ZixTreeIter* ti = NULL;
-  ZixTree*     t  = zix_tree_new(false, int_cmp, NULL, NULL, NULL);
+  ZixTree*     t  = zix_tree_new(NULL, false, int_cmp, NULL, NULL, NULL);
 
   // Insert n_elems elements
   struct timespec insert_start = bench_start();
@@ -157,7 +157,7 @@ bench_zix_btree(size_t n_elems,
 
   uintptr_t    r  = 0u;
   ZixBTreeIter ti = zix_btree_end_iter;
-  ZixBTree*    t  = zix_btree_new(int_cmp, NULL);
+  ZixBTree*    t  = zix_btree_new(NULL, int_cmp, NULL);
 
   // Insert n_elems elements
   struct timespec insert_start = bench_start();
diff --git a/include/zix/allocator.h b/include/zix/allocator.h
new file mode 100644
index 0000000..73e2207
--- /dev/null
+++ b/include/zix/allocator.h
@@ -0,0 +1,152 @@
+/*
+  Copyright 2021 David Robillard <d@drobilla.net>
+
+  Permission to use, copy, modify, and/or distribute this software for any
+  purpose with or without fee is hereby granted, provided that the above
+  copyright notice and this permission notice appear in all copies.
+
+  THIS SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+*/
+
+#ifndef ZIX_ALLOCATOR_H
+#define ZIX_ALLOCATOR_H
+
+#include "zix/attributes.h"
+
+#include <stddef.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+   @addtogroup zix
+   @{
+   @name Allocator
+   @{
+*/
+
+/// Opaque user data for allocation functions
+typedef void ZixAllocatorHandle;
+
+/**
+   General malloc-like memory allocation function.
+
+   This works like the standard C malloc(), except has an additional handle
+   parameter for implementing stateful allocators without static data.
+*/
+typedef void* ZIX_ALLOCATED (
+  *ZixMallocFunc)(ZixAllocatorHandle* ZIX_NULLABLE handle, size_t size);
+
+/**
+   General calloc-like memory allocation function.
+
+   This works like the standard C calloc(), except has an additional handle
+   parameter for implementing stateful allocators without static data.
+*/
+typedef void* ZIX_ALLOCATED (*ZixCallocFunc)(ZixAllocatorHandle* ZIX_NULLABLE
+                                                    handle,
+                                             size_t nmemb,
+                                             size_t size);
+
+/**
+   General realloc-like memory reallocation function.
+
+   This works like the standard C remalloc(), except has an additional handle
+   parameter for implementing stateful allocators without static data.
+*/
+typedef void* ZIX_ALLOCATED (*ZixReallocFunc)(ZixAllocatorHandle* ZIX_NULLABLE
+                                                                 handle,
+                                              void* ZIX_NULLABLE ptr,
+                                              size_t             size);
+
+/**
+   General free-like memory deallocation function.
+
+   This works like the standard C remalloc(), except has an additional handle
+   parameter for implementing stateful allocators without static data.
+*/
+typedef void (*ZixFreeFunc)(ZixAllocatorHandle* ZIX_NULLABLE handle,
+                            void* ZIX_NULLABLE               ptr);
+
+/**
+   A memory allocator.
+
+   This object-like structure provides an interface like the standard C
+   functions malloc(), calloc(), realloc(), and free().  It allows the user to
+   pass a custom allocator to be used by data structures.
+*/
+typedef struct {
+  ZixAllocatorHandle* ZIX_NULLABLE handle;
+  ZixMallocFunc ZIX_NONNULL        malloc;
+  ZixCallocFunc ZIX_NONNULL        calloc;
+  ZixReallocFunc ZIX_NONNULL       realloc;
+  ZixFreeFunc ZIX_NONNULL          free;
+} ZixAllocator;
+
+/// Return the default allocator which simply uses the system allocator
+ZIX_CONST_API
+const ZixAllocator* ZIX_NONNULL
+zix_default_allocator(void);
+
+/// Convenience wrapper that defers to malloc() if allocator is null
+static inline void* ZIX_ALLOCATED
+zix_malloc(const ZixAllocator* const ZIX_NULLABLE allocator, const size_t size)
+{
+  const ZixAllocator* const actual =
+    allocator ? allocator : zix_default_allocator();
+
+  return actual->malloc(actual->handle, size);
+}
+
+/// Convenience wrapper that defers to calloc() if allocator is null
+static inline void* ZIX_ALLOCATED
+zix_calloc(const ZixAllocator* const ZIX_NULLABLE allocator,
+           const size_t                           nmemb,
+           const size_t                           size)
+{
+  const ZixAllocator* const actual =
+    allocator ? allocator : zix_default_allocator();
+
+  return actual->calloc(actual->handle, nmemb, size);
+}
+
+/// Convenience wrapper that defers to realloc() if allocator is null
+static inline void* ZIX_ALLOCATED
+zix_realloc(const ZixAllocator* const ZIX_NULLABLE allocator,
+            void* const ZIX_NULLABLE               ptr,
+            const size_t                           size)
+{
+  const ZixAllocator* const actual =
+    allocator ? allocator : zix_default_allocator();
+
+  return actual->realloc(actual->handle, ptr, size);
+}
+
+/// Convenience wrapper that defers to free() if allocator is null
+static inline void
+zix_free(const ZixAllocator* const ZIX_NULLABLE allocator,
+         void* const ZIX_NULLABLE               ptr)
+{
+  const ZixAllocator* const actual =
+    allocator ? allocator : zix_default_allocator();
+
+  actual->free(actual->handle, ptr);
+}
+
+/**
+   @}
+   @}
+*/
+
+#ifdef __cplusplus
+} /* extern "C" */
+#endif
+
+#endif /* ZIX_ALLOCATOR_H */
diff --git a/include/zix/btree.h b/include/zix/btree.h
index 3979084..9c67dda 100644
--- a/include/zix/btree.h
+++ b/include/zix/btree.h
@@ -17,6 +17,7 @@
 #ifndef ZIX_BTREE_H
 #define ZIX_BTREE_H
 
+#include "zix/allocator.h"
 #include "zix/attributes.h"
 #include "zix/common.h"
 
@@ -88,7 +89,9 @@ static const ZixBTreeIter zix_btree_end_iter = {
 */
 ZIX_API
 ZixBTree* ZIX_ALLOCATED
-zix_btree_new(ZixComparator ZIX_NONNULL cmp, const void* ZIX_NULLABLE cmp_data);
+zix_btree_new(const ZixAllocator* ZIX_NULLABLE allocator,
+              ZixComparator ZIX_NONNULL        cmp,
+              const void* ZIX_NULLABLE         cmp_data);
 
 /**
    Free `t` and all the nodes it contains.
diff --git a/include/zix/hash.h b/include/zix/hash.h
index 8be8941..cc56d0e 100644
--- a/include/zix/hash.h
+++ b/include/zix/hash.h
@@ -17,6 +17,7 @@
 #ifndef ZIX_HASH_H
 #define ZIX_HASH_H
 
+#include "zix/allocator.h"
 #include "zix/attributes.h"
 #include "zix/common.h"
 
@@ -133,9 +134,10 @@ typedef struct {
 */
 ZIX_API
 ZixHash* ZIX_ALLOCATED
-zix_hash_new(ZixKeyFunc ZIX_NONNULL      key_func,
-             ZixHashFunc ZIX_NONNULL     hash_func,
-             ZixKeyEqualFunc ZIX_NONNULL equal_func);
+zix_hash_new(const ZixAllocator* ZIX_NULLABLE allocator,
+             ZixKeyFunc ZIX_NONNULL           key_func,
+             ZixHashFunc ZIX_NONNULL          hash_func,
+             ZixKeyEqualFunc ZIX_NONNULL      equal_func);
 
 /// Free `hash`
 ZIX_API
diff --git a/include/zix/ring.h b/include/zix/ring.h
index b323e70..af98ecf 100644
--- a/include/zix/ring.h
+++ b/include/zix/ring.h
@@ -17,6 +17,7 @@
 #ifndef ZIX_RING_H
 #define ZIX_RING_H
 
+#include "zix/allocator.h"
 #include "zix/attributes.h"
 
 #include <stdint.h>
@@ -48,7 +49,7 @@ typedef struct ZixRingImpl ZixRing;
 */
 ZIX_MALLOC_API
 ZixRing* ZIX_ALLOCATED
-zix_ring_new(uint32_t size);
+zix_ring_new(const ZixAllocator* ZIX_NULLABLE allocator, uint32_t size);
 
 /// Destroy a ring
 ZIX_API
diff --git a/include/zix/tree.h b/include/zix/tree.h
index 2ce7490..c3b9262 100644
--- a/include/zix/tree.h
+++ b/include/zix/tree.h
@@ -17,6 +17,7 @@
 #ifndef ZIX_TREE_H
 #define ZIX_TREE_H
 
+#include "zix/allocator.h"
 #include "zix/attributes.h"
 #include "zix/common.h"
 
@@ -43,11 +44,12 @@ typedef struct ZixTreeNodeImpl ZixTreeIter;
 /// Create a new (empty) tree
 ZIX_API
 ZixTree* ZIX_ALLOCATED
-zix_tree_new(bool                        allow_duplicates,
-             ZixComparator ZIX_NONNULL   cmp,
-             void* ZIX_NULLABLE          cmp_data,
-             ZixDestroyFunc ZIX_NULLABLE destroy,
-             const void* ZIX_NULLABLE    destroy_user_data);
+zix_tree_new(const ZixAllocator* ZIX_NULLABLE allocator,
+             bool                             allow_duplicates,
+             ZixComparator ZIX_NONNULL        cmp,
+             void* ZIX_NULLABLE               cmp_data,
+             ZixDestroyFunc ZIX_NULLABLE      destroy,
+             const void* ZIX_NULLABLE         destroy_user_data);
 
 /// Free `t`
 ZIX_API
diff --git a/meson.build b/meson.build
index f732a5f..2da63a0 100644
--- a/meson.build
+++ b/meson.build
@@ -172,6 +172,7 @@ endif
 #
 
 sources = [
+  'src/allocator.c',
   'src/bitset.c',
   'src/btree.c',
   'src/digest.c',
@@ -226,6 +227,7 @@ install_headers(headers, subdir: versioned_name / 'zix')
 #
 
 tests = [
+  'allocator_test',
   'bitset_test',
   'btree_test',
   'digest_test',
diff --git a/src/allocator.c b/src/allocator.c
new file mode 100644
index 0000000..ccd48ea
--- /dev/null
+++ b/src/allocator.c
@@ -0,0 +1,67 @@
+/*
+  Copyright 2011-2021 David Robillard <d@drobilla.net>
+
+  Permission to use, copy, modify, and/or distribute this software for any
+  purpose with or without fee is hereby granted, provided that the above
+  copyright notice and this permission notice appear in all copies.
+
+  THIS SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+*/
+
+#include "zix/allocator.h"
+
+#include <stdlib.h>
+
+ZIX_MALLOC_FUNC
+static void*
+zix_default_malloc(ZixAllocatorHandle* const handle, const size_t size)
+{
+  (void)handle;
+  return malloc(size);
+}
+
+ZIX_MALLOC_FUNC
+static void*
+zix_default_calloc(ZixAllocatorHandle* const handle,
+                   const size_t              nmemb,
+                   const size_t              size)
+{
+  (void)handle;
+  return calloc(nmemb, size);
+}
+
+static void*
+zix_default_realloc(ZixAllocatorHandle* const handle,
+                    void* const               ptr,
+                    const size_t              size)
+{
+  (void)handle;
+  return realloc(ptr, size);
+}
+
+static void
+zix_default_free(ZixAllocatorHandle* const handle, void* const ptr)
+{
+  (void)handle;
+  free(ptr);
+}
+
+const ZixAllocator*
+zix_default_allocator(void)
+{
+  static const ZixAllocator default_allocator = {
+    NULL,
+    zix_default_malloc,
+    zix_default_calloc,
+    zix_default_realloc,
+    zix_default_free,
+  };
+
+  return &default_allocator;
+}
diff --git a/src/btree.c b/src/btree.c
index 7d967b8..b74eb92 100644
--- a/src/btree.c
+++ b/src/btree.c
@@ -18,7 +18,6 @@
 
 #include <assert.h>
 #include <stdint.h>
-#include <stdlib.h>
 #include <string.h>
 
 // #define ZIX_BTREE_SORTED_CHECK 1
@@ -39,10 +38,11 @@ typedef uint16_t ZixShort;
 #define ZIX_BTREE_INODE_VALS (ZIX_BTREE_LEAF_VALS / 2u)
 
 struct ZixBTreeImpl {
-  ZixBTreeNode* root;
-  ZixComparator cmp;
-  const void*   cmp_data;
-  size_t        size;
+  const ZixAllocator* allocator;
+  ZixBTreeNode*       root;
+  ZixComparator       cmp;
+  const void*         cmp_data;
+  size_t              size;
 };
 
 struct ZixBTreeNodeImpl {
@@ -67,7 +67,7 @@ static_assert(sizeof(ZixBTreeNode) == ZIX_BTREE_PAGE_SIZE, "");
 #endif
 
 static ZixBTreeNode*
-zix_btree_node_new(const bool leaf)
+zix_btree_node_new(const ZixAllocator* const allocator, const bool leaf)
 {
 #if !((defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112l) || \
       (defined(__cplusplus) && __cplusplus >= 201103L))
@@ -76,7 +76,9 @@ zix_btree_node_new(const bool leaf)
          ZIX_BTREE_PAGE_SIZE - 2u * sizeof(ZixBTreeNode*));
 #endif
 
-  ZixBTreeNode* const node = (ZixBTreeNode*)malloc(sizeof(ZixBTreeNode));
+  ZixBTreeNode* const node =
+    (ZixBTreeNode*)zix_malloc(allocator, sizeof(ZixBTreeNode));
+
   if (node) {
     node->is_leaf = leaf;
     node->n_vals  = 0u;
@@ -95,23 +97,26 @@ zix_btree_child(const ZixBTreeNode* const node, const unsigned i)
 }
 
 ZixBTree*
-zix_btree_new(const ZixComparator cmp, const void* const cmp_data)
+zix_btree_new(const ZixAllocator* const allocator,
+              const ZixComparator       cmp,
+              const void* const         cmp_data)
 {
   assert(cmp);
 
-  ZixBTree* const t = (ZixBTree*)malloc(sizeof(ZixBTree));
+  ZixBTree* const t = (ZixBTree*)zix_malloc(allocator, sizeof(ZixBTree));
   if (!t) {
     return NULL;
   }
 
-  if (!(t->root = zix_btree_node_new(true))) {
-    free(t);
+  if (!(t->root = zix_btree_node_new(allocator, true))) {
+    zix_free(allocator, t);
     return NULL;
   }
 
-  t->cmp      = cmp;
-  t->cmp_data = cmp_data;
-  t->size     = 0;
+  t->allocator = allocator;
+  t->cmp       = cmp;
+  t->cmp_data  = cmp_data;
+  t->size      = 0;
 
   return t;
 }
@@ -126,7 +131,7 @@ zix_btree_free_children(ZixBTree* const      t,
     for (ZixShort i = 0; i < n->n_vals + 1u; ++i) {
       zix_btree_free_children(
         t, zix_btree_child(n, i), destroy, destroy_user_data);
-      free(zix_btree_child(n, i));
+      zix_free(t->allocator, zix_btree_child(n, i));
     }
   }
 
@@ -150,8 +155,8 @@ zix_btree_free(ZixBTree* const      t,
 {
   if (t) {
     zix_btree_clear(t, destroy, destroy_user_data);
-    free(t->root);
-    free(t);
+    zix_free(t->allocator, t->root);
+    zix_free(t->allocator, t);
   }
 }
 
@@ -208,9 +213,10 @@ zix_btree_aerase(void** const array, const unsigned n, const unsigned i)
 
 /// Split lhs, the i'th child of `n`, into two nodes
 static ZixBTreeNode*
-zix_btree_split_child(ZixBTreeNode* const n,
-                      const unsigned      i,
-                      ZixBTreeNode* const lhs)
+zix_btree_split_child(const ZixAllocator* const allocator,
+                      ZixBTreeNode* const       n,
+                      const unsigned            i,
+                      ZixBTreeNode* const       lhs)
 {
   assert(lhs->n_vals == zix_btree_max_vals(lhs));
   assert(n->n_vals < ZIX_BTREE_INODE_VALS);
@@ -218,7 +224,7 @@ zix_btree_split_child(ZixBTreeNode* const n,
   assert(zix_btree_child(n, i) == lhs);
 
   const ZixShort max_n_vals = zix_btree_max_vals(lhs);
-  ZixBTreeNode*  rhs        = zix_btree_node_new(lhs->is_leaf);
+  ZixBTreeNode*  rhs        = zix_btree_node_new(allocator, lhs->is_leaf);
   if (!rhs) {
     return NULL;
   }
@@ -407,7 +413,7 @@ zix_btree_is_full(const ZixBTreeNode* const n)
 static ZixStatus
 zix_btree_grow_up(ZixBTree* const t)
 {
-  ZixBTreeNode* const new_root = zix_btree_node_new(false);
+  ZixBTreeNode* const new_root = zix_btree_node_new(t->allocator, false);
   if (!new_root) {
     return ZIX_STATUS_NO_MEM;
   }
@@ -416,7 +422,7 @@ zix_btree_grow_up(ZixBTree* const t)
   new_root->data.inode.children[0] = t->root;
 
   // Split the old root to get two balanced siblings
-  zix_btree_split_child(new_root, 0, t->root);
+  zix_btree_split_child(t->allocator, new_root, 0, t->root);
   t->root = new_root;
 
   return ZIX_STATUS_SUCCESS;
@@ -450,7 +456,9 @@ zix_btree_insert(ZixBTree* const t, void* const e)
     ZixBTreeNode* child = node->data.inode.children[i];
     if (zix_btree_is_full(child)) {
       // The child is full, split it before continuing
-      ZixBTreeNode* const rhs = zix_btree_split_child(node, i, child);
+      ZixBTreeNode* const rhs =
+        zix_btree_split_child(t->allocator, node, i, child);
+
       if (!rhs) {
         return ZIX_STATUS_NO_MEM;
       }
@@ -620,10 +628,10 @@ zix_btree_merge(ZixBTree* const t, ZixBTreeNode* const n, const unsigned i)
     // Root is now empty, replace it with its only child
     assert(n == t->root);
     t->root = lhs;
-    free(n);
+    zix_free(t->allocator, n);
   }
 
-  free(rhs);
+  zix_free(t->allocator, rhs);
   return lhs;
 }
 
diff --git a/src/hash.c b/src/hash.c
index 471ccb7..00ab2c0 100644
--- a/src/hash.c
+++ b/src/hash.c
@@ -19,7 +19,6 @@
 #include <assert.h>
 #include <stdbool.h>
 #include <stdint.h>
-#include <stdlib.h>
 
 typedef struct ZixHashEntry {
   ZixHashCode    hash;  ///< Non-folded hash value
@@ -27,32 +26,35 @@ typedef struct ZixHashEntry {
 } ZixHashEntry;
 
 struct ZixHashImpl {
-  ZixKeyFunc      key_func;   ///< User-provided key accessor
-  ZixHashFunc     hash_func;  ///< User-provided hashing function
-  ZixKeyEqualFunc equal_func; ///< User-provided equality comparison function
-  size_t          count;      ///< Number of records stored in the table
-  size_t          mask;       ///< Bit mask for fast modulo (n_entries - 1)
-  size_t          n_entries;  ///< Power of two table size
-  ZixHashEntry*   entries;    ///< Pointer to dynamically allocated table
+  const ZixAllocator* allocator;  ///< User allocator
+  ZixKeyFunc          key_func;   ///< User key accessor
+  ZixHashFunc         hash_func;  ///< User hashing function
+  ZixKeyEqualFunc     equal_func; ///< User equality comparison function
+  size_t              count;      ///< Number of records stored in the table
+  size_t              mask;       ///< Bit mask for fast modulo (n_entries - 1)
+  size_t              n_entries;  ///< Power of two table size
+  ZixHashEntry*       entries;    ///< Pointer to dynamically allocated table
 };
 
 static const size_t min_n_entries = 4u;
 static const size_t tombstone     = 0xDEADu;
 
 ZixHash*
-zix_hash_new(const ZixKeyFunc      key_func,
-             const ZixHashFunc     hash_func,
-             const ZixKeyEqualFunc equal_func)
+zix_hash_new(const ZixAllocator* const allocator,
+             const ZixKeyFunc          key_func,
+             const ZixHashFunc         hash_func,
+             const ZixKeyEqualFunc     equal_func)
 {
   assert(key_func);
   assert(hash_func);
   assert(equal_func);
 
-  ZixHash* const hash = (ZixHash*)malloc(sizeof(ZixHash));
+  ZixHash* const hash = (ZixHash*)zix_malloc(allocator, sizeof(ZixHash));
   if (!hash) {
     return NULL;
   }
 
+  hash->allocator  = allocator;
   hash->key_func   = key_func;
   hash->hash_func  = hash_func;
   hash->equal_func = equal_func;
@@ -60,9 +62,11 @@ zix_hash_new(const ZixKeyFunc      key_func,
   hash->n_entries  = min_n_entries;
   hash->mask       = hash->n_entries - 1u;
 
-  hash->entries = (ZixHashEntry*)calloc(hash->n_entries, sizeof(ZixHashEntry));
+  hash->entries =
+    (ZixHashEntry*)zix_calloc(allocator, hash->n_entries, sizeof(ZixHashEntry));
+
   if (!hash->entries) {
-    free(hash);
+    zix_free(allocator, hash);
     return NULL;
   }
 
@@ -73,8 +77,8 @@ void
 zix_hash_free(ZixHash* const hash)
 {
   if (hash) {
-    free(hash->entries);
-    free(hash);
+    zix_free(hash->allocator, hash->entries);
+    zix_free(hash->allocator, hash);
   }
 }
 
@@ -173,7 +177,9 @@ rehash(ZixHash* const hash, const size_t old_n_entries)
   const size_t        new_n_entries = hash->n_entries;
 
   // Replace the array in the hash first so we can use find_entry() normally
-  hash->entries = (ZixHashEntry*)calloc(new_n_entries, sizeof(ZixHashEntry));
+  hash->entries = (ZixHashEntry*)zix_calloc(
+    hash->allocator, new_n_entries, sizeof(ZixHashEntry));
+
   if (!hash->entries) {
     return ZIX_STATUS_NO_MEM;
   }
@@ -191,7 +197,7 @@ rehash(ZixHash* const hash, const size_t old_n_entries)
     }
   }
 
-  free(old_entries);
+  zix_free(hash->allocator, old_entries);
   return ZIX_STATUS_SUCCESS;
 }
 
diff --git a/src/ring.c b/src/ring.c
index 4a4692f..ed1f32e 100644
--- a/src/ring.c
+++ b/src/ring.c
@@ -44,11 +44,12 @@
 #endif
 
 struct ZixRingImpl {
-  uint32_t write_head; ///< Read index into buf
-  uint32_t read_head;  ///< Write index into buf
-  uint32_t size;       ///< Size (capacity) in bytes
-  uint32_t size_mask;  ///< Mask for fast modulo
-  char*    buf;        ///< Contents
+  const ZixAllocator* allocator;  ///< User allocator
+  uint32_t            write_head; ///< Read index into buf
+  uint32_t            read_head;  ///< Write index into buf
+  uint32_t            size;       ///< Size (capacity) in bytes
+  uint32_t            size_mask;  ///< Mask for fast modulo
+  char*               buf;        ///< Contents
 };
 
 static inline uint32_t
@@ -66,14 +67,23 @@ next_power_of_two(uint32_t size)
 }
 
 ZixRing*
-zix_ring_new(uint32_t size)
+zix_ring_new(const ZixAllocator* const allocator, uint32_t size)
 {
-  ZixRing* ring    = (ZixRing*)malloc(sizeof(ZixRing));
-  ring->write_head = 0;
-  ring->read_head  = 0;
-  ring->size       = next_power_of_two(size);
-  ring->size_mask  = ring->size - 1;
-  ring->buf        = (char*)malloc(ring->size);
+  ZixRing* ring = (ZixRing*)zix_malloc(allocator, sizeof(ZixRing));
+
+  if (ring) {
+    ring->allocator  = allocator;
+    ring->write_head = 0;
+    ring->read_head  = 0;
+    ring->size       = next_power_of_two(size);
+    ring->size_mask  = ring->size - 1;
+
+    if (!(ring->buf = (char*)zix_malloc(allocator, ring->size))) {
+      zix_free(allocator, ring);
+      return NULL;
+    }
+  }
+
   return ring;
 }
 
@@ -81,8 +91,8 @@ void
 zix_ring_free(ZixRing* ring)
 {
   if (ring) {
-    free(ring->buf);
-    free(ring);
+    zix_free(ring->allocator, ring->buf);
+    zix_free(ring->allocator, ring);
   }
 }
 
diff --git a/src/tree.c b/src/tree.c
index f26ff90..94cbab5 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -19,19 +19,19 @@
 #include "zix/common.h"
 
 #include <assert.h>
-#include <stdlib.h>
 #include <string.h>
 
 typedef struct ZixTreeNodeImpl ZixTreeNode;
 
 struct ZixTreeImpl {
-  ZixTreeNode*   root;
-  ZixDestroyFunc destroy;
-  const void*    destroy_user_data;
-  ZixComparator  cmp;
-  void*          cmp_data;
-  size_t         size;
-  bool           allow_duplicates;
+  const ZixAllocator* allocator;
+  ZixTreeNode*        root;
+  ZixDestroyFunc      destroy;
+  const void*         destroy_user_data;
+  ZixComparator       cmp;
+  void*               cmp_data;
+  size_t              size;
+  bool                allow_duplicates;
 };
 
 struct ZixTreeNodeImpl {
@@ -67,20 +67,26 @@ struct ZixTreeNodeImpl {
 #endif
 
 ZixTree*
-zix_tree_new(bool           allow_duplicates,
-             ZixComparator  cmp,
-             void*          cmp_data,
-             ZixDestroyFunc destroy,
-             const void*    destroy_user_data)
+zix_tree_new(const ZixAllocator* const allocator,
+             bool                      allow_duplicates,
+             ZixComparator             cmp,
+             void*                     cmp_data,
+             ZixDestroyFunc            destroy,
+             const void*               destroy_user_data)
 {
-  ZixTree* t           = (ZixTree*)malloc(sizeof(ZixTree));
-  t->root              = NULL;
-  t->destroy           = destroy;
-  t->destroy_user_data = destroy_user_data;
-  t->cmp               = cmp;
-  t->cmp_data          = cmp_data;
-  t->size              = 0;
-  t->allow_duplicates  = allow_duplicates;
+  ZixTree* t = (ZixTree*)zix_malloc(allocator, sizeof(ZixTree));
+
+  if (t) {
+    t->allocator         = allocator;
+    t->root              = NULL;
+    t->destroy           = destroy;
+    t->destroy_user_data = destroy_user_data;
+    t->cmp               = cmp;
+    t->cmp_data          = cmp_data;
+    t->size              = 0;
+    t->allow_duplicates  = allow_duplicates;
+  }
+
   return t;
 }
 
@@ -93,7 +99,8 @@ zix_tree_free_rec(ZixTree* t, ZixTreeNode* n)
     if (t->destroy) {
       t->destroy(n->data, t->destroy_user_data);
     }
-    free(n);
+
+    zix_free(t->allocator, n);
   }
 }
 
@@ -102,7 +109,7 @@ zix_tree_free(ZixTree* t)
 {
   if (t) {
     zix_tree_free_rec(t, t->root);
-    free(t);
+    zix_free(t->allocator, t);
   }
 }
 
@@ -370,7 +377,7 @@ zix_tree_insert(ZixTree* t, void* e, ZixTreeIter** ti)
   }
 
   // Allocate a new node n
-  if (!(n = (ZixTreeNode*)malloc(sizeof(ZixTreeNode)))) {
+  if (!(n = (ZixTreeNode*)zix_malloc(t->allocator, sizeof(ZixTreeNode)))) {
     return ZIX_STATUS_NO_MEM;
   }
   memset(n, '\0', sizeof(ZixTreeNode));
@@ -456,7 +463,7 @@ zix_tree_remove(ZixTree* t, ZixTreeIter* ti)
     if (t->destroy) {
       t->destroy(n->data, t->destroy_user_data);
     }
-    free(n);
+    zix_free(t->allocator, n);
     --t->size;
     assert(t->size == 0);
     return ZIX_STATUS_SUCCESS;
@@ -588,7 +595,7 @@ zix_tree_remove(ZixTree* t, ZixTreeIter* ti)
   if (t->destroy) {
     t->destroy(n->data, t->destroy_user_data);
   }
-  free(n);
+  zix_free(t->allocator, n);
 
   --t->size;
 
diff --git a/test/allocator_test.c b/test/allocator_test.c
new file mode 100644
index 0000000..432b836
--- /dev/null
+++ b/test/allocator_test.c
@@ -0,0 +1,65 @@
+/*
+  Copyright 2014-2021 David Robillard <d@drobilla.net>
+
+  Permission to use, copy, modify, and/or distribute this software for any
+  purpose with or without fee is hereby granted, provided that the above
+  copyright notice and this permission notice appear in all copies.
+
+  THIS SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+*/
+
+#undef NDEBUG
+
+#include "zix/allocator.h"
+
+#include <assert.h>
+
+static void
+test_allocator(void)
+{
+  // Just a basic smoke test to check that things seem to be working
+
+  const ZixAllocator* const allocator = zix_default_allocator();
+
+  char* const malloced = (char*)zix_malloc(allocator, 4);
+  malloced[0]          = 0;
+  malloced[3]          = 3;
+  assert(malloced[0] == 0);
+  assert(malloced[3] == 3);
+
+  char* const calloced = (char*)zix_calloc(allocator, 4, 1);
+  assert(calloced[0] == 0);
+  assert(calloced[1] == 0);
+  assert(calloced[2] == 0);
+  assert(calloced[3] == 0);
+
+  char* const realloced = (char*)zix_realloc(allocator, calloced, 8);
+  assert(realloced[0] == 0);
+  assert(realloced[1] == 0);
+  assert(realloced[2] == 0);
+  assert(realloced[3] == 0);
+  realloced[4] = 4;
+  realloced[5] = 5;
+  realloced[6] = 6;
+  realloced[7] = 7;
+  assert(realloced[4] == 4);
+  assert(realloced[5] == 5);
+  assert(realloced[6] == 6);
+  assert(realloced[7] == 7);
+
+  zix_free(allocator, realloced);
+  zix_free(allocator, malloced);
+}
+
+int
+main(void)
+{
+  test_allocator();
+  return 0;
+}
diff --git a/test/btree_test.c b/test/btree_test.c
index a890f0f..00a5790 100644
--- a/test/btree_test.c
+++ b/test/btree_test.c
@@ -136,7 +136,7 @@ no_destroy(void* const ptr, const void* const user_data)
 static void
 test_clear(void)
 {
-  ZixBTree* t = zix_btree_new(int_cmp, NULL);
+  ZixBTree* t = zix_btree_new(NULL, int_cmp, NULL);
 
   for (uintptr_t r = 0u; r < n_clear_insertions; ++r) {
     assert(!zix_btree_insert(t, (void*)(r + 1u)));
@@ -151,7 +151,7 @@ test_clear(void)
 static void
 test_free(void)
 {
-  ZixBTree* t = zix_btree_new(int_cmp, NULL);
+  ZixBTree* t = zix_btree_new(NULL, int_cmp, NULL);
 
   for (uintptr_t r = 0u; r < n_clear_insertions; ++r) {
     assert(!zix_btree_insert(t, (void*)(r + 1u)));
@@ -167,7 +167,7 @@ test_iter_comparison(void)
 {
   static const size_t n_elems = 4096u;
 
-  ZixBTree* const t = zix_btree_new(int_cmp, NULL);
+  ZixBTree* const t = zix_btree_new(NULL, int_cmp, NULL);
 
   // Store increasing numbers from 1 (jammed into the pointers themselves)
   for (uintptr_t r = 1u; r < n_elems; ++r) {
@@ -211,7 +211,7 @@ test_insert_split_value(void)
   static const size_t    n_insertions = 767u; // Number of insertions to split
   static const uintptr_t split_value  = 512u; // Value that will be pulled up
 
-  ZixBTree* const t = zix_btree_new(int_cmp, NULL);
+  ZixBTree* const t = zix_btree_new(NULL, int_cmp, NULL);
 
   // Insert right up until it would cause a split
   for (uintptr_t r = 1u; r < n_insertions; ++r) {
@@ -235,7 +235,7 @@ test_remove_cases(void)
   static const uintptr_t s2           = 255u;
   static const size_t    n_insertions = s1 * s2 * 1000u;
 
-  ZixBTree* const t = zix_btree_new(int_cmp, NULL);
+  ZixBTree* const t = zix_btree_new(NULL, int_cmp, NULL);
 
   // Insert in s1-sized chunks
   for (uintptr_t phase = 0u; phase < s1; ++phase) {
@@ -270,7 +270,7 @@ stress(const unsigned test_num, const size_t n_elems)
   }
 
   uintptr_t r  = 0;
-  ZixBTree* t  = zix_btree_new(int_cmp, NULL);
+  ZixBTree* t  = zix_btree_new(NULL, int_cmp, NULL);
   ZixStatus st = ZIX_STATUS_SUCCESS;
 
   if (!t) {
@@ -554,7 +554,7 @@ stress(const unsigned test_num, const size_t n_elems)
   // Test lower_bound with wildcard comparator
 
   TestContext ctx = {test_num, n_elems};
-  if (!(t = zix_btree_new(wildcard_cmp, &ctx))) {
+  if (!(t = zix_btree_new(NULL, wildcard_cmp, &ctx))) {
     return test_fail(t, "Failed to allocate tree\n");
   }
 
diff --git a/test/hash_test.c b/test/hash_test.c
index 834c391..62e2d18 100644
--- a/test/hash_test.c
+++ b/test/hash_test.c
@@ -132,7 +132,7 @@ string_equal(const char* const a, const char* const b)
 static int
 stress_with(const ZixHashFunc hash_func, const size_t n_elems)
 {
-  ZixHash* hash = zix_hash_new(identity, hash_func, string_equal);
+  ZixHash* hash = zix_hash_new(NULL, identity, hash_func, string_equal);
   if (!hash) {
     return test_fail("Failed to allocate hash\n");
   }
@@ -336,8 +336,9 @@ test_all_tombstones(void)
     "3 b",
   };
 
-  ZixStatus st   = ZIX_STATUS_SUCCESS;
-  ZixHash*  hash = zix_hash_new(identity, identity_index_hash, string_equal);
+  ZixStatus st = ZIX_STATUS_SUCCESS;
+  ZixHash*  hash =
+    zix_hash_new(NULL, identity, identity_index_hash, string_equal);
 
   // Insert each element then immediately remove it
   for (unsigned i = 0u; i < N_STRINGS; ++i) {
diff --git a/test/ring_test.c b/test/ring_test.c
index 61a6d1f..23c5adc 100644
--- a/test/ring_test.c
+++ b/test/ring_test.c
@@ -137,7 +137,7 @@ main(int argc, char** argv)
          MSG_SIZE,
          size);
 
-  ring = zix_ring_new(size);
+  ring = zix_ring_new(NULL, size);
   assert(ring);
   if (zix_ring_read_space(ring) != 0) {
     return failure("New ring is not empty\n");
diff --git a/test/tree_test.c b/test/tree_test.c
index 76c247d..7d2b6af 100644
--- a/test/tree_test.c
+++ b/test/tree_test.c
@@ -63,7 +63,7 @@ stress(unsigned test_num, size_t n_elems)
 {
   uintptr_t    r  = 0u;
   ZixTreeIter* ti = NULL;
-  ZixTree*     t  = zix_tree_new(true, int_cmp, NULL, NULL, NULL);
+  ZixTree*     t  = zix_tree_new(NULL, true, int_cmp, NULL, NULL, NULL);
 
   // Insert n_elems elements
   for (size_t i = 0; i < n_elems; ++i) {