6 files changed, 99 insertions, 32 deletions

diff --git a/mem/include/mem.h b/mem/include/mem.h
index 892ea4f..224b069 100644
--- a/mem/include/mem.h
+++ b/mem/include/mem.h
@@ -86,9 +86,15 @@
 #define mem_get_chunk_handle(MEM, CHUNK_PTR) \
   mem_get_chunk_handle_(&(MEM)->mem, CHUNK_PTR)
 
-/// Return the total capacity of the allocator in bytes.
+/// Return the block size in bytes.
+#define mem_block_size_bytes(MEM) memp_block_size_bytes_(&(MEM)->pool)
+
+/// Return the total capacity of the allocator.
 #define mem_capacity(MEM) mem_capacity_(&(MEM)->mem)
 
+/// Return the number of used blocks in the allocator.
+#define mem_size(MEM) mem_size_(&(MEM)->mem)
+
 /// Set whether to trap when attempting to allocate beyond capacity.
 #define mem_enable_traps(MEM, enable) mem_enable_traps_(&(MEM)->mem, enable)
 
@@ -97,18 +103,20 @@
 /// The caller can use 'i' as the index of the current chunk.
 ///
 /// It is valid to mem_free() the chunk at each step of the iteration.
-#define mem_foreach(MEM, ITER, BODY)                                  \
-  {                                                                   \
-    size_t i = 0;                                                     \
-    do {                                                              \
-      if ((MEM)->mem.chunks[i].used) {                                \
-        __typeof__((MEM)->object[0])* ITER =                          \
-            &(((__typeof__((MEM)->object[0])*)(MEM)->mem.blocks))[i]; \
-        (void)ITER;                                                   \
-        BODY;                                                         \
-      }                                                               \
-      i = (MEM)->mem.chunks[i].next;                                  \
-    } while (i);                                                      \
+#define mem_foreach(MEM, ITER, BODY)                                    \
+  {                                                                     \
+    size_t i = 0;                                                       \
+    if ((MEM)->mem.num_used_blocks > 0) {                               \
+      do {                                                              \
+        if ((MEM)->mem.chunks[i].used) {                                \
+          __typeof__((MEM)->object[0])* ITER =                          \
+              &(((__typeof__((MEM)->object[0])*)(MEM)->mem.blocks))[i]; \
+          (void)ITER;                                                   \
+          BODY;                                                         \
+        }                                                               \
+        i = (MEM)->mem.chunks[i].next;                                  \
+      } while (i);                                                      \
+    }                                                                   \
   }
 
 // -----------------------------------------------------------------------------
@@ -137,6 +145,7 @@ typedef struct Chunk {
 typedef struct Memory {
   size_t   block_size_bytes;
   size_t   num_blocks;
+  size_t   num_used_blocks;
   size_t   next_free_chunk;
   bool     dynamic; /// True if blocks and chunks are dynamically-allocated.
   bool     trap;    /// Whether to trap when allocating beyond capacity.
@@ -164,5 +173,7 @@ void*  mem_alloc_(Memory*, size_t num_blocks);
 void   mem_free_(Memory*, void** chunk_ptr);
 void*  mem_get_chunk_(const Memory*, size_t chunk_handle);
 size_t mem_get_chunk_handle_(const Memory*, const void* chunk);
+size_t mem_block_size_bytes_(const Memory*);
 size_t mem_capacity_(const Memory*);
+size_t mem_size_(const Memory*);
 void   mem_enable_traps_(Memory*, bool);
diff --git a/mem/src/mem.c b/mem/src/mem.c
index 2904035..4f5e5ef 100644
--- a/mem/src/mem.c
+++ b/mem/src/mem.c
@@ -14,6 +14,7 @@ bool mem_make_(
 
   mem->block_size_bytes = block_size_bytes;
   mem->num_blocks       = num_blocks;
+  mem->num_used_blocks  = 0;
   mem->next_free_chunk  = 0;
   mem->trap             = true;
 
@@ -107,6 +108,7 @@ void* mem_alloc_(Memory* mem, size_t num_blocks) {
   } while (chunk_idx != start);
 
   if (found) {
+    mem->num_used_blocks++;
     mem->next_free_chunk = mem->chunks[chunk_idx].next;
     return &mem->blocks[chunk_idx * mem->block_size_bytes];
   } else {
@@ -168,6 +170,8 @@ void mem_free_(Memory* mem, void** chunk_ptr) {
     }
   }
 
+  mem->num_used_blocks--;
+
   *chunk_ptr = 0;
 }
 
@@ -189,9 +193,19 @@ size_t mem_get_chunk_handle_(const Memory* mem, const void* chunk) {
   return block_byte_index / mem->block_size_bytes;
 }
 
+size_t mem_block_size_bytes_(const Memory* mem) {
+  assert(mem);
+  return mem->block_size_bytes;
+}
+
 size_t mem_capacity_(const Memory* mem) {
   assert(mem);
-  return mem->num_blocks * mem->block_size_bytes;
+  return mem->num_blocks;
+}
+
+size_t mem_size_(const Memory* mem) {
+  assert(mem);
+  return mem->num_used_blocks;
 }
 
 void mem_enable_traps_(Memory* mem, bool enable) {
diff --git a/mem/test/mem_test.c b/mem/test/mem_test.c
index d3c43b9..97db079 100644
--- a/mem/test/mem_test.c
+++ b/mem/test/mem_test.c
@@ -41,6 +41,8 @@ TEST_CASE(mem_fully_allocate) {
     const int* block = mem_alloc(&mem, 1);
     TEST_TRUE(block != 0);
   }
+
+  TEST_TRUE(mem_size(&mem) == NUM_BLOCKS);
 }
 
 // Allocate N chunks of 1 block each, then free them.
@@ -60,6 +62,7 @@ TEST_CASE(mem_fill_then_free) {
   }
 
   TEST_EQUAL(count(&mem), 0);
+  TEST_TRUE(mem_size(&mem) == 0);
 }
 
 // Attempt to allocate blocks past the maximum allocator size.
@@ -78,6 +81,8 @@ TEST_CASE(mem_allocate_beyond_max_size) {
   for (int i = 0; i < NUM_BLOCKS; ++i) {
     TEST_EQUAL(mem_alloc(&mem, 1), 0);
   }
+
+  TEST_TRUE(mem_size(&mem) == NUM_BLOCKS);
 }
 
 // Free blocks should always remain zeroed out.
@@ -114,6 +119,7 @@ TEST_CASE(mem_traverse_empty) {
   mem_make(&mem);
 
   TEST_EQUAL(count(&mem), 0);
+  TEST_TRUE(mem_size(&mem) == 0);
 }
 
 // Traverse a partially full allocator.
@@ -130,6 +136,7 @@ TEST_CASE(mem_traverse_partially_full) {
   }
 
   TEST_EQUAL(sum(&mem), (N) * (N + 1) / 2);
+  TEST_TRUE(mem_size(&mem) == N);
 }
 
 // Traverse a full allocator.
@@ -144,6 +151,7 @@ TEST_CASE(mem_traverse_full) {
   }
 
   TEST_EQUAL(sum(&mem), (NUM_BLOCKS) * (NUM_BLOCKS + 1) / 2);
+  TEST_TRUE(mem_size(&mem) == NUM_BLOCKS);
 }
 
 // Get the ith (allocated) chunk.
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.