3 files changed, 47 insertions, 14 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.