3 files changed, 52 insertions, 18 deletions

diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
index de9ea4f..245173b 100644
--- a/mempool/include/mempool.h
+++ b/mempool/include/mempool.h
@@ -86,9 +86,15 @@
 #define mempool_get_block_index(POOL, BLOCK_PTR) \
   mempool_get_block_index_(&(POOL)->pool, BLOCK_PTR)
 
-/// Return the total capacity of the mempool in bytes.
+/// Return the block size in bytes.
+#define mempool_block_size_bytes(POOL) mempool_block_size_bytes_(&(POOL)->pool)
+
+/// Return the total capacity of the mempool.
 #define mempool_capacity(POOL) mempool_capacity_(&(POOL)->pool)
 
+/// Return the number of used blocks in the mempool.
+#define mempool_size(POOL) mempool_size_(&(POOL)->pool)
+
 /// Set whether to trap when attempting to allocate beyond capacity.
 #define mempool_enable_traps(POOL, enable) \
   mempool_enable_traps_(&(POOL)->pool, enable)
@@ -98,22 +104,24 @@
 /// The caller can use 'i' as the index of the current block.
 ///
 /// It is valid to mempool_free() the object at each step of the iteration.
-#define mempool_foreach(POOL, ITER, BODY)                                \
-  {                                                                      \
-    size_t i = (POOL)->pool.used;                                        \
-    do {                                                                 \
-      if ((POOL)->pool.block_info[i].used) {                             \
-        __typeof__((POOL)->object[0])* ITER =                            \
-            &(((__typeof__((POOL)->object[0])*)(POOL)->pool.blocks))[i]; \
-        (void)ITER;                                                      \
-        BODY;                                                            \
-      }                                                                  \
-      const size_t next = (POOL)->pool.block_info[i].next_used;          \
-      if (next == i) {                                                   \
-        break;                                                           \
-      }                                                                  \
-      i = next;                                                          \
-    } while (true);                                                      \
+#define mempool_foreach(POOL, ITER, BODY)                                  \
+  {                                                                        \
+    size_t i = (POOL)->pool.used;                                          \
+    if ((POOL)->pool.num_used_blocks > 0) {                                \
+      do {                                                                 \
+        if ((POOL)->pool.block_info[i].used) {                             \
+          __typeof__((POOL)->object[0])* ITER =                            \
+              &(((__typeof__((POOL)->object[0])*)(POOL)->pool.blocks))[i]; \
+          (void)ITER;                                                      \
+          BODY;                                                            \
+        }                                                                  \
+        const size_t next = (POOL)->pool.block_info[i].next_used;          \
+        if (next == i) {                                                   \
+          break;                                                           \
+        }                                                                  \
+        i = next;                                                          \
+      } while (true);                                                      \
+    }                                                                      \
   }
 
 // -----------------------------------------------------------------------------
@@ -133,6 +141,7 @@ typedef struct BlockInfo {
 typedef struct mempool {
   size_t     block_size_bytes;
   size_t     num_blocks;
+  size_t     num_used_blocks;
   size_t     head;    /// Points to the first block in the free list.
   size_t     used;    /// Points to the first block in the used list.
   bool       dynamic; /// True if blocks and info are dynamically-allocated.
@@ -161,5 +170,7 @@ void*  mempool_alloc_(mempool*);
 void   mempool_free_(mempool*, void** block_ptr);
 void*  mempool_get_block_(const mempool*, size_t block_index);
 size_t mempool_get_block_index_(const mempool*, const void* block);
+size_t mempool_block_size_bytes_(const mempool*);
 size_t mempool_capacity_(const mempool*);
+size_t mempool_size_(const mempool*);
 void   mempool_enable_traps_(mempool*, bool);
diff --git a/mempool/src/mempool.c b/mempool/src/mempool.c
index b09038b..46f1053 100644
--- a/mempool/src/mempool.c
+++ b/mempool/src/mempool.c
@@ -24,6 +24,7 @@ bool mempool_make_(
 
   pool->block_size_bytes = block_size_bytes;
   pool->num_blocks       = num_blocks;
+  pool->num_used_blocks  = 0;
   pool->head             = 0;
   pool->used             = 0;
   pool->trap             = true;
@@ -91,6 +92,8 @@ void* mempool_alloc_(mempool* pool) {
   pool->head      = head->next_free;
   head->next_free = 0;
 
+  pool->num_used_blocks++;
+
   return block;
 }
 
@@ -119,6 +122,8 @@ void mempool_free_(mempool* pool, void** block_ptr) {
     pool->used = 0;
   }
 
+  pool->num_used_blocks--;
+
   *block_ptr = 0;
 }
 
@@ -136,9 +141,19 @@ size_t mempool_get_block_index_(const mempool* pool, const void* block) {
   return block_byte_index / pool->block_size_bytes;
 }
 
+size_t mempool_block_size_bytes_(const mempool* pool) {
+  assert(pool);
+  return pool->block_size_bytes;
+}
+
 size_t mempool_capacity_(const mempool* pool) {
   assert(pool);
-  return pool->num_blocks * pool->block_size_bytes;
+  return pool->num_blocks;
+}
+
+size_t mempool_size_(const mempool* pool) {
+  assert(pool);
+  return pool->num_used_blocks;
 }
 
 void mempool_enable_traps_(mempool* pool, bool enable) {
diff --git a/mempool/test/mempool_test.c b/mempool/test/mempool_test.c
index 6c48a2a..5eaee51 100644
--- a/mempool/test/mempool_test.c
+++ b/mempool/test/mempool_test.c
@@ -41,6 +41,8 @@ TEST_CASE(mempool_allocate_until_full) {
     const int* block = mempool_alloc(&pool);
     TEST_TRUE(block != 0);
   }
+
+  TEST_TRUE(mempool_size(&pool) == NUM_BLOCKS);
 }
 
 // Allocate all N blocks, then free them.
@@ -60,6 +62,7 @@ TEST_CASE(mempool_fill_then_free) {
   }
 
   TEST_EQUAL(count(&pool), 0);
+  TEST_TRUE(mempool_size(&pool) == 0);
 }
 
 // Attempt to allocate blocks past the maximum pool size.
@@ -78,6 +81,8 @@ TEST_CASE(mempool_allocate_beyond_max_size) {
   for (int i = 0; i < NUM_BLOCKS; ++i) {
     TEST_EQUAL(mempool_alloc(&pool), 0);
   }
+
+  TEST_TRUE(mempool_size(&pool) == NUM_BLOCKS);
 }
 
 // Free blocks should always remain zeroed out.
@@ -114,6 +119,7 @@ TEST_CASE(mempool_traverse_empty) {
   mempool_make(&pool);
 
   TEST_EQUAL(count(&pool), 0);
+  TEST_TRUE(mempool_size(&pool) == 0);
 }
 
 // Traverse a partially full pool.
@@ -130,6 +136,7 @@ TEST_CASE(mempool_traverse_partially_full) {
   }
 
   TEST_EQUAL(sum(&pool), N * (N + 1) / 2);
+  TEST_TRUE(mempool_size(&pool) == N);
 }
 
 // Traverse a full pool.
@@ -144,6 +151,7 @@ TEST_CASE(mempool_traverse_full) {
   }
 
   TEST_EQUAL(sum(&pool), NUM_BLOCKS * (NUM_BLOCKS + 1) / 2);
+  TEST_TRUE(mempool_size(&pool) == NUM_BLOCKS);
 }
 
 // Get the ith (allocated) block.