1 files changed, 61 insertions, 28 deletions

diff --git a/hicn-light/src/hicn/base/pool.c b/hicn-light/src/hicn/base/pool.c
index 0f5f728e3..31abb13f1 100644
--- a/hicn-light/src/hicn/base/pool.c
+++ b/hicn-light/src/hicn/base/pool.c
@@ -30,36 +30,52 @@
 #include "common.h"
 #include "pool.h"
 
+#include <stdio.h> // XXX
+
 void
-_pool_init(void ** pool_ptr, size_t elt_size, size_t max_size)
+_pool_init(void ** pool_ptr, size_t elt_size, size_t init_size, size_t max_size)
 {
     assert(pool_ptr);
     assert(elt_size);
 
-    pool_hdr_t * ph = calloc(POOL_HDRLEN + elt_size * max_size, 1);
-    if (!ph) {
-        *pool_ptr = NULL;
-        return;
-    }
+    init_size = next_pow2(init_size);
+
+    if (max_size && init_size > max_size)
+        goto ERR_MAX_SIZE;
+
+    /* The initial pool size is rounded to the next power of two */
+    size_t alloc_size = next_pow2(init_size);
+
+    pool_hdr_t * ph = calloc(POOL_HDRLEN + alloc_size * elt_size, 1);
+    if (!ph)
+        goto ERR_MALLOC;
 
     ph->elt_size = elt_size;
+    ph->alloc_size = alloc_size;
     ph->max_size = max_size;
 
     /* Free indices */
     off_t * free_indices;
-    vector_init(free_indices, max_size);
-    for(unsigned i = 0; i < max_size; i++)
-        free_indices[i] = (max_size - 1) - i;
-    vector_len(free_indices) = max_size;
+    vector_init(free_indices, init_size, max_size);
+    for(unsigned i = 0; i < init_size; i++)
+        free_indices[i] = (init_size - 1) - i;
+    vector_len(free_indices) = init_size;
     ph->free_indices = free_indices;
 
     /* Free bitmap */
     uint_fast32_t * fb = ph->free_bitmap;
-    bitmap_init(fb, max_size);
-    bitmap_set_to(fb, max_size);
+    bitmap_init(fb, init_size, max_size);
+    bitmap_set_to(fb, init_size);
     ph->free_bitmap = fb;
 
     *pool_ptr = (uint8_t*)ph + POOL_HDRLEN;
+
+    return;
+
+ERR_MALLOC:
+ERR_MAX_SIZE:
+    *pool_ptr = NULL;
+    return;
 }
 
 void
@@ -73,42 +89,59 @@ void
 _pool_resize(void ** pool_ptr, size_t elt_size)
 {
     pool_hdr_t * ph = pool_hdr(*pool_ptr);
-    size_t old_elts = ph->max_size;
-    size_t new_elts = old_elts * 2;
+    size_t old_size = ph->alloc_size;
+    size_t new_size = old_size * 2;
+
+    if (ph->max_size && new_size > ph->max_size)
+        goto ERR_MAX_SIZE;
 
     /* Double pool storage */
-    ph = realloc(ph, POOL_HDRLEN + new_elts * elt_size);
-    if (!ph) {
-        *pool_ptr = NULL;
-        return;
-    }
+    ph = realloc(ph, POOL_HDRLEN + new_size * elt_size);
+    if (!ph)
+        goto ERR_REALLOC;
     ph->elt_size = elt_size;
-    ph->max_size = new_elts;
+    ph->alloc_size = new_size;
 
     /*
-     * After resize, the pool will have old_elts free indices, ranging from
-     * old_elts to (new_elts - 1)
+     * After resize, the pool will have new free indices, ranging from
+     * old_size to (new_size - 1)
      */
-    off_t * free_indices = ph->free_indices;
-    vector_ensure_pos(free_indices, old_elts);
-    for (unsigned i = 0; i < old_elts; i++)
-        free_indices[i] = new_elts - 1 - i;
+    vector_ensure_pos(ph->free_indices, old_size);
+    for (unsigned i = 0; i < old_size; i++)
+        ph->free_indices[i] = new_size - 1 - i;
+    vector_len(ph->free_indices) = old_size;
+
+    /* We also need to update the bitmap */
+    bitmap_ensure_pos(ph->free_bitmap, new_size - 1);
+    bitmap_set_range(ph->free_bitmap, old_size, new_size - 1);
 
     /* Reassign pool pointer */
     *pool_ptr = (uint8_t*)ph + POOL_HDRLEN;
+
+    return;
+
+ERR_REALLOC:
+ERR_MAX_SIZE:
+    *pool_ptr = NULL;
+    return;
 }
 
-void
+off_t
 _pool_get(void ** pool_ptr, void ** elt, size_t elt_size)
 {
     pool_hdr_t * ph = pool_hdr(*pool_ptr);
     uint64_t l = vector_len(ph->free_indices);
-    if (l == 0)
+    if (l == 0) {
         _pool_resize(pool_ptr, elt_size);
+        ph = pool_hdr(*pool_ptr);
+        l = vector_len(ph->free_indices);
+    }
     off_t free_id = ph->free_indices[l - 1];
     vector_len(ph->free_indices)--;
+    bitmap_unset(ph->free_bitmap, free_id);
     *elt = *pool_ptr + free_id;
     memset(*elt, 0, sizeof(elt));
+    return free_id;
 }
 
 void