[Qemu-devel] [PATCH 14/20] hbitmap: add hbitmap_alloc_with_data and hbitmap_required_size

[Paolo Bonzini <pbonzini@redhat.com>]

Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>
---
 hbitmap.c | 57 +++++++++++++++++++++++++++++++++++++++++++++++----------
 hbitmap.h | 25 +++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/hbitmap.c b/hbitmap.c
index ae59f39..aa4586f 100644
--- a/hbitmap.c
+++ b/hbitmap.c
@@ -81,6 +81,9 @@ struct HBitmap {
      */
     int granularity;
 
+    /* True if the last level should be freed by hbitmap_free.  */
+    bool last_level_allocated;
+
     /* A number of progressively less coarse bitmaps (i.e. level 0 is the
      * coarsest).  Each bit in level N represents a word in level N+1 that
      * has a set bit, except the last level where each bit represents the
@@ -367,34 +370,68 @@ bool hbitmap_get(const HBitmap *hb, uint64_t item)
 
 void hbitmap_free(HBitmap *hb)
 {
-    unsigned i;
-    for (i = HBITMAP_LEVELS; i-- > 0; ) {
+    unsigned i = HBITMAP_LEVELS;
+
+    if (!hb->last_level_allocated) {
+        i--;
+    }
+
+    while (i-- > 0) {
         g_free(hb->levels[i]);
     }
     g_free(hb);
 }
 
-HBitmap *hbitmap_alloc(uint64_t size, int granularity)
+static uint64_t hbitmap_round_size(uint64_t size, int granularity)
 {
-    HBitmap *hb = g_malloc0(sizeof(struct HBitmap));
-    unsigned i;
-
     assert(granularity >= 0 && granularity < 64);
+
+    if (!size) {
+        size = 1;
+    }
+
     size = (size + (1ULL << granularity) - 1) >> granularity;
     assert(size <= ((uint64_t)1 << HBITMAP_LOG_MAX_SIZE));
+    return size;
+}
 
-    hb->size = size;
+size_t hbitmap_required_size(uint64_t size, int granularity)
+{
+    size_t longs;
+
+    size = hbitmap_round_size(size, granularity);
+    longs = (size + BITS_PER_LONG - 1) >> BITS_PER_LEVEL;
+    return longs * sizeof(unsigned long);
+}
+
+HBitmap *hbitmap_alloc_with_data(uint64_t size, int granularity, void *data)
+{
+    HBitmap *hb = g_malloc0(sizeof(struct HBitmap));
+    unsigned i;
+
+    hb->size = hbitmap_round_size(size, granularity);
     hb->granularity = granularity;
+    hb->last_level_allocated = (data == NULL);
+
     for (i = HBITMAP_LEVELS; i-- > 0; ) {
-        size = MAX((size + BITS_PER_LONG - 1) >> BITS_PER_LEVEL, 1);
-        hb->levels[i] = g_malloc0(size * sizeof(unsigned long));
+        if (data == NULL) {
+            data = g_malloc0(hbitmap_required_size(size, granularity));
+        }
+        hb->levels[i] = data;
+        data = NULL;
+        granularity += BITS_PER_LEVEL;
     }
 
     /* We necessarily have free bits in level 0 due to the definition
      * of HBITMAP_LEVELS, so use one for a sentinel.  This speeds up
      * hbitmap_iter_skip_words.
      */
-    assert(size == 1);
+    assert(hbitmap_required_size(size, granularity) == sizeof(unsigned long));
     hb->levels[0][0] |= 1UL << (BITS_PER_LONG - 1);
     return hb;
 }
+
+HBitmap *hbitmap_alloc(uint64_t size, int granularity)
+{
+    return hbitmap_alloc_with_data(size, granularity, NULL);
+}
diff --git a/hbitmap.h b/hbitmap.h
index 7ddfb66..35ee51a 100644
--- a/hbitmap.h
+++ b/hbitmap.h
@@ -64,6 +64,31 @@ struct HBitmapIter {
 HBitmap *hbitmap_alloc(uint64_t size, int granularity);
 
 /**
+ * hbitmap_required_size:
+ * @size: Number of bits in the bitmap.
+ * @granularity: Granularity of the bitmap.
+ *
+ * Return the number of bytes that are needed for a bitmap with the
+ * given size and granularity.
+ *
+ * A block of this size can then be passed to @hbitmap_alloc_with_data.
+ */
+size_t hbitmap_required_size(uint64_t size, int granularity);
+
+/**
+ * hbitmap_alloc_with_data:
+ * @size: Number of bits in the bitmap.
+ * @granularity: Granularity of the bitmap.
+ * @data: Pointer to a data block that will be used for the bottom level
+ * of the HBitmap.
+ *
+ * Allocate a new HBitmap, using a client-provided data block for the
+ * actual bitmap and allocating memory only for the compressed levels.
+ * If @data is NULL, this function is equivalent to @hbitmap_alloc.
+ */
+HBitmap *hbitmap_alloc_with_data(uint64_t size, int granularity, void *data);
+
+/**
  * hbitmap_empty:
  * @hb: HBitmap to operate on.
  *
-- 
1.8.0.1