Fix mem and mempool iteration.

1 files changed, 23 insertions, 21 deletions

diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
index 8251a70..23786b3 100644
--- a/mempool/include/mempool.h
+++ b/mempool/include/mempool.h
@@ -19,19 +19,22 @@
 #include <stdint.h>
 
 /// Define a statically-allocated, typed pool of the given number of blocks.
-#define DEF_MEMPOOL(POOL, TYPE, NUM_BLOCKS) \
-  typedef struct POOL {                     \
-    mempool   pool;                         \
-    BlockInfo block_info[NUM_BLOCKS];       \
-    TYPE      blocks[NUM_BLOCKS];           \
+#define DEF_MEMPOOL(POOL, TYPE, NUM_BLOCKS)                   \
+  typedef struct POOL {                                       \
+    mempool   pool;                                           \
+    BlockInfo block_info[NUM_BLOCKS];                         \
+    TYPE      blocks[NUM_BLOCKS];                             \
+    /* For uniformity with the dynamically-allocated pool. */ \
+    TYPE* object;                                             \
   } POOL;
 
 /// Define a dynamically-allocated, typed pool.
-#define DEF_MEMPOOL_DYN(POOL, TYPE) \
-  typedef struct POOL {             \
-    mempool    pool;                \
-    BlockInfo* block_info;          \
-    TYPE*      blocks;              \
+#define DEF_MEMPOOL_DYN(POOL, TYPE)                                          \
+  typedef struct POOL {                                                      \
+    mempool pool;                                                            \
+    /* Does not point anywhere. Storing a pointer here so that we can recall \
+     * the type. */                                                          \
+    TYPE* object;                                                            \
   } POOL;
 
 /// Initialize a statically-allocated pool.
@@ -74,7 +77,7 @@
 /// Return the ith block.
 /// The block must have been allocated.
 #define mempool_get_block(POOL, INDEX) \
-  ((__typeof__((POOL)->blocks[0])*)mempool_get_block_(&(POOL)->pool, INDEX))
+  ((__typeof__((POOL)->object[0])*)mempool_get_block_(&(POOL)->pool, INDEX))
 
 /// Get the index to the given block.
 #define mempool_get_block_index(POOL, BLOCK_PTR) \
@@ -88,16 +91,15 @@
 /// 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 < (sizeof((POOL)->blocks) / sizeof(__typeof__((POOL)->blocks[0]))); \
-       ++i) {                                                                \
-    if (!(POOL)->block_info[i].used) {                                       \
-      continue;                                                              \
-    }                                                                        \
-    __typeof__((POOL)->blocks[0])* ITER = &(POOL)->blocks[i];                \
-    (void)ITER;                                                              \
-    BODY;                                                                    \
+#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;                                                            \
   }
 
 // -----------------------------------------------------------------------------