Fix mem and mempool iteration.

2 files changed, 28 insertions, 22 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);
 
 // -----------------------------------------------------------------------------
diff --git a/mem/test/mem_test.c b/mem/test/mem_test.c
index 6ab4c7c..2f242c3 100644
--- a/mem/test/mem_test.c
+++ b/mem/test/mem_test.c
@@ -4,7 +4,7 @@
 
 #define NUM_BLOCKS 10
 
-DEF_MEM(test_mem, int, NUM_BLOCKS);
+DEF_MEM(test_mem, int, NUM_BLOCKS)
 
 static int count(test_mem* mem) {
   int count = 0;
@@ -27,7 +27,7 @@ TEST_CASE(mem_create) {
 // Create a dynamically-backed allocator.
 TEST_CASE(mem_create_dyn) {
   DEF_MEM_DYN(dyn_mem, int);
-  
+
   dyn_mem mem;
   mem_make_dyn(&mem, NUM_BLOCKS, sizeof(int));
 }