2 files changed, 43 insertions, 41 deletions

diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
index 994e25a..b76ae7c 100644
--- a/mempool/include/mempool.h
+++ b/mempool/include/mempool.h
@@ -100,14 +100,14 @@
 // -----------------------------------------------------------------------------
 
 typedef struct BlockInfo {
-  bool used;
+  size_t next; /// For free blocks, points to the next free block.
+  bool   used;
 } BlockInfo;
 
 typedef struct mempool {
   size_t     block_size_bytes;
   size_t     num_blocks;
-  size_t     next_free_block;
-  bool       full;
+  size_t     head;    /// Points to the first block in the free list.
   bool       dynamic; /// True if blocks and info are dynamically-allocated.
   BlockInfo* block_info;
   uint8_t*   blocks;
diff --git a/mempool/src/mempool.c b/mempool/src/mempool.c
index 059db93..fdb5f8c 100644
--- a/mempool/src/mempool.c
+++ b/mempool/src/mempool.c
@@ -3,22 +3,39 @@
 #include <stdlib.h>
 #include <string.h>
 
+#define NO_BLOCK ((size_t)-1)
+
 static inline size_t min(size_t a, size_t b) { return a < b ? a : b; }
 
+/// Initialize the free list.
+/// All of the blocks in the pool are assumed free.
+static void init_free_list(mempool* pool) {
+  assert(pool);
+  for (size_t i = 0; i < pool->num_blocks - 1; ++i) {
+    pool->block_info[i].next = i + 1;
+  }
+  pool->block_info[pool->num_blocks - 1].next = NO_BLOCK;
+}
+
 bool mempool_make_(
     mempool* pool, BlockInfo* block_info, void* blocks, size_t num_blocks,
     size_t block_size_bytes) {
   assert(pool);
   assert((block_info && blocks) || (!block_info && !blocks));
   assert(num_blocks >= 1);
+
   pool->block_size_bytes = block_size_bytes;
   pool->num_blocks       = num_blocks;
-  pool->next_free_block  = 0;
-  pool->full             = false;
+  pool->head             = 0;
+
+  // Initialize blocks and block info.
   if (!block_info) {
     block_info    = calloc(num_blocks, sizeof(BlockInfo));
     blocks        = calloc(num_blocks, block_size_bytes);
     pool->dynamic = true;
+    if ((block_info == 0) || (blocks == 0)) {
+      return false;
+    }
   } else {
     memset(blocks, 0, num_blocks * block_size_bytes);
     memset(block_info, 0, num_blocks * sizeof(BlockInfo));
@@ -26,7 +43,10 @@ bool mempool_make_(
   }
   pool->block_info = block_info;
   pool->blocks     = blocks;
-  return (block_info != 0) && (blocks != 0);
+
+  init_free_list(pool);
+
+  return true;
 }
 
 void mempool_del_(mempool* pool) {
@@ -45,37 +65,24 @@ void mempool_del_(mempool* pool) {
 
 void mempool_clear_(mempool* pool) {
   assert(pool);
-  pool->next_free_block = 0;
-  pool->full            = false;
+  pool->head = 0;
   memset(pool->blocks, 0, pool->num_blocks * pool->block_size_bytes);
   memset(pool->block_info, 0, pool->num_blocks * sizeof(BlockInfo));
+  init_free_list(pool);
 }
 
 void* mempool_alloc_(mempool* pool) {
   assert(pool);
 
-  if (pool->full) {
+  if (pool->head == NO_BLOCK) {
     return 0;
   }
 
   // Allocate the block.
-  void* block = &pool->blocks[pool->next_free_block * pool->block_size_bytes];
-  pool->block_info[pool->next_free_block].used = true;
-
-  // Search for the next free block. If it does not exist, flag the pool
-  // full.
-  //
-  // The search starts near the current free block, where we might be more
-  // likely to find free blocks. The search starts with i=1 since we only
-  // need to test the other N-1 blocks in the pool.
-  pool->full = true;
-  for (size_t i = 1; i < pool->num_blocks; i++) {
-    pool->next_free_block = (pool->next_free_block + 1) % pool->num_blocks;
-    if (!pool->block_info[pool->next_free_block].used) {
-      pool->full = false;
-      break;
-    }
-  }
+  BlockInfo* head  = &pool->block_info[pool->head];
+  void*      block = &pool->blocks[pool->head * pool->block_size_bytes];
+  head->used       = true;
+  pool->head       = head->next;
 
   return block;
 }
@@ -84,27 +91,22 @@ void mempool_free_(mempool* pool, void** block_ptr) {
   assert(pool);
   assert(block_ptr);
 
-  // Zero out the block so that we don't get stray values the next time it is
-  // allocated.
-  memset(*block_ptr, 0, pool->block_size_bytes);
-
   const size_t block_index =
       ((uint8_t*)*block_ptr - pool->blocks) / pool->block_size_bytes;
   assert(block_index < pool->num_blocks);
+  BlockInfo* info = &pool->block_info[block_index];
 
   // Disallow double-frees.
-  assert(pool->block_info[block_index].used);
+  assert(info->used);
 
-  pool->block_info[block_index].used = false;
-  if (pool->full) {
-    pool->next_free_block = block_index;
-    pool->full            = false;
-  } else {
-    // Prefer to allocate blocks towards the start of the pool. This way, blocks
-    // should cluster around this area and the pool should offer better memory
-    // locality for used blocks.
-    pool->next_free_block = min(pool->next_free_block, block_index);
-  }
+  // Zero out the block so that we don't get stray values the next time it is
+  // allocated.
+  memset(*block_ptr, 0, pool->block_size_bytes);
+
+  // Free the block and add it to the head of the free list.
+  info->used = false;
+  info->next = pool->head;
+  pool->head = block_index;
 
   *block_ptr = 0;
 }