Add used list to mempool; fix mem iteration.

3 files changed, 58 insertions, 36 deletions

diff --git a/mem/include/mem.h b/mem/include/mem.h
index b3d9157..bcff39f 100644
--- a/mem/include/mem.h
+++ b/mem/include/mem.h
@@ -92,17 +92,19 @@
 /// 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)->chunks[i].next;                                      \
-  } while (i);
+#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);                                                      \
+  }
 
 // -----------------------------------------------------------------------------
 
diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
index 23786b3..bd4d4dd 100644
--- a/mempool/include/mempool.h
+++ b/mempool/include/mempool.h
@@ -91,28 +91,43 @@
 /// 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)                            \
-  for (size_t i = 0; i < (POOL)->pool.num_blocks; ++i) {             \
-    if (!(POOL)->pool.block_info[i].used) {                          \
-      continue;                                                      \
-    }                                                                \
-    __typeof__((POOL)->object[0])* ITER =                            \
-        &(((__typeof__((POOL)->object[0])*)(POOL)->pool.blocks))[i]; \
-    (void)ITER;                                                      \
-    BODY;                                                            \
+#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);                                                      \
   }
 
 // -----------------------------------------------------------------------------
 
 typedef struct BlockInfo {
-  size_t next; /// For free blocks, points to the next free block.
+  size_t next_free; /// For free blocks, points to the next free block.
+  size_t next_used; /// For used blocks, points to the next used block.
   bool   used;
 } BlockInfo;
 
+/// Memory pool.
+///
+/// 'head' and 'used' always points to a valid block (e.g., 0).
+/// The implementation must check whether the head of the lists are used/free.
+/// For example, iteration must stop when it finds the first unused block
+/// (BlockInfo.used == 0).
 typedef struct mempool {
   size_t     block_size_bytes;
   size_t     num_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.
   BlockInfo* block_info;
   uint8_t*   blocks;
diff --git a/mempool/src/mempool.c b/mempool/src/mempool.c
index 679f124..1100dad 100644
--- a/mempool/src/mempool.c
+++ b/mempool/src/mempool.c
@@ -3,18 +3,14 @@
 #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[i].next_free = i + 1;
   }
-  pool->block_info[pool->num_blocks - 1].next = NO_BLOCK;
+  pool->block_info[pool->num_blocks - 1].next_free = 0;
 }
 
 bool mempool_make_(
@@ -27,6 +23,7 @@ bool mempool_make_(
   pool->block_size_bytes = block_size_bytes;
   pool->num_blocks       = num_blocks;
   pool->head             = 0;
+  pool->used             = 0;
 
   // Initialize blocks and block info.
   if (!block_info) {
@@ -66,6 +63,7 @@ void mempool_del_(mempool* pool) {
 void mempool_clear_(mempool* pool) {
   assert(pool);
   pool->head = 0;
+  pool->used = 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);
@@ -74,15 +72,18 @@ void mempool_clear_(mempool* pool) {
 void* mempool_alloc_(mempool* pool) {
   assert(pool);
 
-  if (pool->head == NO_BLOCK) {
-    return 0;
+  BlockInfo* head = &pool->block_info[pool->head];
+  if (head->used) {
+    return 0; // Pool is full.
   }
 
   // Allocate the block.
-  BlockInfo* head  = &pool->block_info[pool->head];
-  void*      block = &pool->blocks[pool->head * pool->block_size_bytes];
-  head->used       = true;
-  pool->head       = head->next;
+  void* block     = &pool->blocks[pool->head * pool->block_size_bytes];
+  head->used      = true;
+  head->next_used = pool->used;
+  pool->used      = pool->head;
+  pool->head      = head->next_free;
+  head->next_free = 0;
 
   return block;
 }
@@ -104,9 +105,13 @@ void mempool_free_(mempool* pool, void** block_ptr) {
   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;
+  info->used      = false;
+  info->next_used = 0;
+  info->next_free = pool->head;
+  pool->head      = block_index;
+  if (pool->used == block_index) {
+    pool->used = 0;
+  }
 
   *block_ptr = 0;
 }