1 files changed, 25 insertions, 10 deletions

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;