Fix mem and mempool iteration.

1 files changed, 26 insertions, 20 deletions

diff --git a/mem/include/mem.h b/mem/include/mem.h
index 69f426f..b3d9157 100644
--- a/mem/include/mem.h
+++ b/mem/include/mem.h
@@ -19,19 +19,24 @@
 #include <stdint.h>
 
 /// Define a typed memory allocator backed by a statically-allocated array.
-#define DEF_MEM(MEM, TYPE, NUM_BLOCKS) \
-  typedef struct MEM {                 \
-    Memory mem;                        \
-    Chunk  chunks[NUM_BLOCKS];         \
-    TYPE   blocks[NUM_BLOCKS];         \
+#define DEF_MEM(MEM, TYPE, NUM_BLOCKS)                        \
+  typedef struct MEM {                                        \
+    Memory mem;                                               \
+    Chunk  chunks[NUM_BLOCKS];                                \
+    TYPE   blocks[NUM_BLOCKS];                                \
+    /* For uniformity with the dynamically-allocated pool. */ \
+    TYPE* object;                                             \
   } MEM;
 
 /// Define a typed memory allocator backed by a dynamically-allocated array.
-#define DEF_MEM_DYN(MEM, TYPE) \
-  typedef struct MEM {         \
-    Memory mem;                \
-    Chunk* chunks;             \
-    TYPE*  blocks;             \
+#define DEF_MEM_DYN(MEM, TYPE)                                               \
+  typedef struct MEM {                                                       \
+    Memory mem;                                                              \
+    Chunk* chunks;                                                           \
+    TYPE*  blocks;                                                           \
+    /* Does not point anywhere. Storing a pointer here so that we can recall \
+     * the type. */                                                          \
+    TYPE* object;                                                            \
   } MEM;
 
 /// Initialize a statically-backed memory allocator.
@@ -73,7 +78,7 @@
 /// Return a pointer to a chunk given the chunk's handle.
 /// The chunk must have been allocated.
 #define mem_get_chunk(MEM, HANDLE) \
-  ((__typeof__((MEM)->blocks[0])*)mem_get_chunk_(&(MEM)->mem, HANDLE))
+  ((__typeof__((MEM)->object[0])*)mem_get_chunk_(&(MEM)->mem, HANDLE))
 
 /// Get the handle to the given chunk.
 #define mem_get_chunk_handle(MEM, CHUNK_PTR) \
@@ -87,15 +92,16 @@
 /// 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)->chunks[i].used) {                              \
-      __typeof__((MEM)->blocks[0])* ITER = &(MEM)->blocks[i]; \
-      (void)ITER;                                             \
-      BODY;                                                   \
-    }                                                         \
-    i = (MEM)->chunks[i].next;                                \
+#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);
 
 // -----------------------------------------------------------------------------