Update listpool's interface to be on par with mempool's.

3 files changed, 56 insertions, 6 deletions

diff --git a/listpool/include/listpool.h b/listpool/include/listpool.h
index 1711449..85a3b27 100644
--- a/listpool/include/listpool.h
+++ b/listpool/include/listpool.h
@@ -45,13 +45,26 @@
     });                             \
   }
 
+/// Return the ith block.
+/// The block must have been allocated.
+#define listpool_get_block(POOL, INDEX) \
+  ((typeof((POOL)->blocks[0])*)listpool_get_block_(&(POOL)->pool, INDEX))
+
+/// Get the index to the given block.
+#define listpool_get_block_index(POOL, BLOCK_PTR) \
+  listpool_get_block_index_(&(POOL)->pool, BLOCK_PTR)
+
 /// Iterate over the used items of the pool.
-#define listpool_foreach(POOL, ITER, BODY)                    \
-  for (list* it_ = (POOL)->pool.used; it_; it_ = it_->next) { \
-    typeof((POOL)->blocks[0])* ITER =                         \
-        &(POOL)->blocks[it_ - (POOL)->pool.nodes];            \
-    (void)ITER;                                               \
-    BODY;                                                     \
+///
+/// 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 listpool_foreach(POOL, ITER, BODY)                      \
+  for (list* it_ = (POOL)->pool.used; it_; it_ = it_->next) {   \
+    const size_t               i    = it_ - (POOL)->pool.nodes; \
+    typeof((POOL)->blocks[0])* ITER = &(POOL)->blocks[i];       \
+    (void)ITER;                                                 \
+    BODY;                                                       \
   }
 
 typedef struct listpool {
@@ -78,3 +91,10 @@ void* listpool_alloc_(listpool* pool);
 /// Free the block.
 /// The block pointer is conveniently set to 0.
 void listpool_free_(listpool* pool, void** block_ptr);
+
+/// Return the ith block.
+/// The block must have been allocated.
+void* listpool_get_block_(const listpool*, size_t block_index);
+
+/// Get the index to the given block.
+size_t listpool_get_block_index_(const listpool*, const void* block);
diff --git a/listpool/src/listpool.c b/listpool/src/listpool.c
index 8e49f32..758062c 100644
--- a/listpool/src/listpool.c
+++ b/listpool/src/listpool.c
@@ -77,3 +77,16 @@ void listpool_free_(listpool* pool, void** block_ptr) {
 
   *block_ptr = 0;
 }
+
+void* listpool_get_block_(const listpool* pool, size_t block_index) {
+  assert(pool);
+  assert(block_index < pool->num_blocks);
+  return pool->blocks + block_index * pool->block_size_bytes;
+}
+
+size_t listpool_get_block_index_(const listpool* pool, const void* block) {
+  assert(pool);
+  const size_t block_byte_index = (const uint8_t*)block - pool->blocks;
+  assert(block_byte_index % pool->block_size_bytes == 0);
+  return block_byte_index / pool->block_size_bytes;
+}
diff --git a/listpool/test/listpool_test.c b/listpool/test/listpool_test.c
index cb54d00..7a7b0cf 100644
--- a/listpool/test/listpool_test.c
+++ b/listpool/test/listpool_test.c
@@ -137,6 +137,23 @@ TEST_CASE(listpool_traverse_full) {
   TEST_EQUAL(sum(&pool), (NUM_BLOCKS) * (NUM_BLOCKS + 1) / 2);
 }
 
+// Get the ith (allocated) block.
+TEST_CASE(listpool_get_block) {
+  test_pool pool;
+  listpool_make(&pool);
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    int* block = listpool_alloc(&pool);
+    TEST_TRUE(block != 0);
+    *block = i;
+    TEST_EQUAL(listpool_get_block_index(&pool, block), (size_t)i);
+  }
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    TEST_EQUAL(*listpool_get_block(&pool, i), i);
+  }
+}
+
 // Remove a value from the list.
 TEST_CASE(listpool_remove_value) {
   test_pool pool;