2 files changed, 84 insertions, 35 deletions

diff --git a/libdimension/pool.c b/libdimension/pool.c
index 82cd46f..a7d6b4f 100644
--- a/libdimension/pool.c
+++ b/libdimension/pool.c
@@ -26,51 +26,63 @@
 #include "dimension-internal.h"
 #include <stdatomic.h>
 
-/// A single allocation and associated destructor.
-typedef struct dmnsn_allocation {
-  void *ptr;
-  dmnsn_callback_fn *cleanup_fn;
-} dmnsn_allocation;
-
 /// Number of pointers per block, we want a block to fit in a single page.
-#define DMNSN_POOL_BLOCK_SIZE ((4096 - 4*sizeof(void *))/sizeof(dmnsn_allocation))
+#define DMNSN_POOL_BLOCK_SIZE (4096/sizeof(void *) - 4)
+/// Number of pointers per tidy block
+#define DMNSN_TIDY_BLOCK_SIZE ((4096 - 4*sizeof(void *))/(sizeof(void *) + sizeof(dmnsn_callback_fn *)))
 
-/// A single block in a thread pool.
+/// A single block in a thread-specific pool.
 typedef struct dmnsn_pool_block {
   /// Current index into allocs[].
   size_t i;
   /// All allocations in the current block.
-  dmnsn_allocation allocs[DMNSN_POOL_BLOCK_SIZE];
+  void *allocs[DMNSN_POOL_BLOCK_SIZE];
   /// Tail pointer to the previous block in the global chain.
   struct dmnsn_pool_block *prev;
 } dmnsn_pool_block;
 
+/// A single tidy block in a thread-specific pool.
+typedef struct dmnsn_tidy_block {
+  /// Current index into allocs[].
+  size_t i;
+  /// All allocations in the current block.
+  void *allocs[DMNSN_TIDY_BLOCK_SIZE];
+  /// All cleanup callbacks in the current block.
+  dmnsn_callback_fn *cleanup_fns[DMNSN_TIDY_BLOCK_SIZE];
+  /// Tail pointer to the previous tidy block in the global chain.
+  struct dmnsn_tidy_block *prev;
+} dmnsn_tidy_block;
+
 /// dmnsn_pool implementation.
 struct dmnsn_pool {
-  /// Thread-local block.
+  /// Thread-local regular block.
   pthread_key_t thread_block;
-  /// Global chain of pools.
+  /// Thread-local tidy block.
+  pthread_key_t thread_tidy_block;
+
+  /// Global chain of regular blocks.
   atomic(dmnsn_pool_block *) chain;
+  /// Global chain of tidy blocks.
+  atomic(dmnsn_tidy_block *) tidy_chain;
 };
 
 dmnsn_pool *
 dmnsn_new_pool(void)
 {
   dmnsn_pool *pool = DMNSN_MALLOC(dmnsn_pool);
+
   dmnsn_key_create(&pool->thread_block, NULL);
+  dmnsn_key_create(&pool->thread_tidy_block, NULL);
+
   atomic_store_explicit(&pool->chain, NULL, memory_order_relaxed);
+  atomic_store_explicit(&pool->tidy_chain, NULL, memory_order_relaxed);
+
   return pool;
 }
 
 void *
 dmnsn_palloc(dmnsn_pool *pool, size_t size)
 {
-  return dmnsn_palloc_tidy(pool, size, NULL);
-}
-
-void *
-dmnsn_palloc_tidy(dmnsn_pool *pool, size_t size, dmnsn_callback_fn *cleanup_fn)
-{
   dmnsn_pool_block *old_block = pthread_getspecific(pool->thread_block);
 
   dmnsn_pool_block *new_block = old_block;
@@ -79,20 +91,46 @@ dmnsn_palloc_tidy(dmnsn_pool *pool, size_t size, dmnsn_callback_fn *cleanup_fn)
     new_block->i = 0;
   }
 
-  dmnsn_allocation *alloc = new_block->allocs + new_block->i;
-  void *result = alloc->ptr = dmnsn_malloc(size);
-  alloc->cleanup_fn = cleanup_fn;
-  ++new_block->i;
+  void *result = dmnsn_malloc(size);
+  new_block->allocs[new_block->i++] = result;
 
   if (dmnsn_unlikely(new_block != old_block)) {
     dmnsn_setspecific(pool->thread_block, new_block);
 
     // Atomically update pool->chain
-    dmnsn_pool_block *chain;
-    do {
-      chain = atomic_load(&pool->chain);
-    } while (!atomic_compare_exchange_weak(&pool->chain, &chain, new_block));
-    new_block->prev = chain;
+    dmnsn_pool_block *old_chain = atomic_exchange(&pool->chain, new_block);
+    new_block->prev = old_chain;
+  }
+
+  return result;
+}
+
+void *
+dmnsn_palloc_tidy(dmnsn_pool *pool, size_t size, dmnsn_callback_fn *cleanup_fn)
+{
+  dmnsn_assert(cleanup_fn != NULL, "NULL cleanup_fn");
+
+  dmnsn_tidy_block *old_block = pthread_getspecific(pool->thread_tidy_block);
+
+  dmnsn_tidy_block *new_block = old_block;
+  if (dmnsn_unlikely(!old_block || old_block->i == DMNSN_TIDY_BLOCK_SIZE)) {
+    new_block = DMNSN_MALLOC(dmnsn_tidy_block);
+    new_block->i = 0;
+  }
+
+  void *result = dmnsn_malloc(size);
+
+  size_t i = new_block->i;
+  new_block->allocs[i] = result;
+  new_block->cleanup_fns[i] = cleanup_fn;
+  ++new_block->i;
+
+  if (dmnsn_unlikely(new_block != old_block)) {
+    dmnsn_setspecific(pool->thread_tidy_block, new_block);
+
+    // Atomically update pool->tidy_chain
+    dmnsn_tidy_block *old_chain = atomic_exchange(&pool->tidy_chain, new_block);
+    new_block->prev = old_chain;
   }
 
   return result;
@@ -107,13 +145,9 @@ dmnsn_delete_pool(dmnsn_pool *pool)
 
   dmnsn_pool_block *block = atomic_load_explicit(&pool->chain, memory_order_relaxed);
   while (block) {
-    // Free all the allocations in reverse order
+    // Free all the allocations
     for (size_t i = block->i; i-- > 0;) {
-      dmnsn_allocation *alloc = block->allocs + i;
-      if (alloc->cleanup_fn) {
-        alloc->cleanup_fn(alloc->ptr);
-      }
-      dmnsn_free(alloc->ptr);
+      dmnsn_free(block->allocs[i]);
     }
 
     // Free the block itself and go to the previous one
@@ -122,6 +156,21 @@ dmnsn_delete_pool(dmnsn_pool *pool)
     dmnsn_free(saved);
   }
 
-  dmnsn_key_delete(pool->thread_block);
+  dmnsn_tidy_block *tidy_block = atomic_load_explicit(&pool->tidy_chain, memory_order_relaxed);
+  while (tidy_block) {
+    // Free all the allocations
+    for (size_t i = tidy_block->i; i-- > 0;) {
+      void *ptr = tidy_block->allocs[i];
+      tidy_block->cleanup_fns[i](ptr);
+      dmnsn_free(ptr);
+    }
+
+    // Free the block itself and go to the previous one
+    dmnsn_tidy_block *saved = tidy_block;
+    tidy_block = tidy_block->prev;
+    dmnsn_free(saved);
+  }
+
+  dmnsn_key_delete(pool->thread_tidy_block);
   dmnsn_free(pool);
 }
diff --git a/libdimension/tests/pool.c b/libdimension/tests/pool.c
index 83f25f7..1722d66 100644
--- a/libdimension/tests/pool.c
+++ b/libdimension/tests/pool.c
@@ -59,10 +59,10 @@ callback(void *ptr)
 
 DMNSN_TEST(pool, callback)
 {
-  DMNSN_PALLOC_TIDY(pool, int, NULL);
+  DMNSN_PALLOC(pool, int);
   DMNSN_PALLOC_TIDY(pool, int, callback);
   DMNSN_PALLOC_TIDY(pool, int, callback);
-  DMNSN_PALLOC_TIDY(pool, int, NULL);
+  DMNSN_PALLOC(pool, int);
 
   dmnsn_delete_pool(pool);
   pool = NULL;