minor optimizations to roaring64_bitmap_or

include/roaring/art/art.h

@@ -199,6 +199,21 @@ void art_iterator_insert(art_iterator_t *iterator, const art_key_chunk_t *key,
  */
 bool art_iterator_erase(art_iterator_t *iterator, art_val_t *erased_val);
 
+/**
+ * Ensures the ART has room for at least `num_leaves` leaves in total before
+ * further leaf allocations are needed. Only leaf storage is reserved; inner
+ * nodes keep growing on demand. Invalidates pointers to values previously
+ * returned by `art_insert` and `art_find`.
+ */
+void art_reserve(art_t *art, uint64_t num_leaves);
+
+/**
+ * Returns the number of leaves (key/value pairs) in the ART. This walks the
+ * free list of leaves, so it is O(1) for an ART that has had no erasures or
+ * that has been shrunken, and O(number of free leaf slots) otherwise.
+ */
+uint64_t art_num_leaves(const art_t *art);
+
 /**
  * Shrinks the internal arrays in the ART to remove any unused elements. Returns
  * the number of bytes freed.

src/art/art.c

@@ -1271,6 +1271,29 @@ static uint8_t art_common_prefix(const art_key_chunk_t key1[],
     return offset;
 }
 
+/**
+ * Grows the array of nodes of the given typecode to `new_capacity` elements,
+ * linking the new elements at the end of the free list. Does nothing if
+ * `new_capacity` is not larger than the current capacity. Invalidates pointers
+ * into the array obtained by `art_deref`.
+ */
+static void art_grow_to(art_t *art, art_typecode_t typecode,
+                        uint64_t new_capacity) {
+    uint64_t capacity = art->capacities[typecode];
+    if (new_capacity <= capacity) {
+        return;
+    }
+    art->capacities[typecode] = new_capacity;
+    art->nodes[typecode] = roaring_realloc(
+        art->nodes[typecode], new_capacity * ART_NODE_SIZES[typecode]);
+    uint64_t increase = new_capacity - capacity;
+    memset(art_get_node(art, capacity, typecode), 0,
+           increase * ART_NODE_SIZES[typecode]);
+    for (uint64_t i = capacity; i < new_capacity; ++i) {
+        art_node_set_next_free(art_get_node(art, i, typecode), typecode, i + 1);
+    }
+}
+
 /**
  * Extends the array of nodes of the given typecode. Invalidates pointers into
  * the array obtained by `art_deref`.
@@ -1289,15 +1312,7 @@ static void art_extend(art_t *art, art_typecode_t typecode) {
     } else {
         new_capacity = 5 * capacity / 4;
     }
-    art->capacities[typecode] = new_capacity;
-    art->nodes[typecode] = roaring_realloc(
-        art->nodes[typecode], new_capacity * ART_NODE_SIZES[typecode]);
-    uint64_t increase = new_capacity - capacity;
-    memset(art_get_node(art, capacity, typecode), 0,
-           increase * ART_NODE_SIZES[typecode]);
-    for (uint64_t i = capacity; i < new_capacity; ++i) {
-        art_node_set_next_free(art_get_node(art, i, typecode), typecode, i + 1);
-    }
+    art_grow_to(art, typecode, new_capacity);
 }
 
 /**
@@ -1846,6 +1861,21 @@ void art_free(art_t *art) {
     }
 }
 
+void art_reserve(art_t *art, uint64_t num_leaves) {
+    art_grow_to(art, CROARING_ART_LEAF_TYPE, num_leaves);
+}
+
+uint64_t art_num_leaves(const art_t *art) {
+    uint64_t num_free = 0;
+    for (uint64_t i = art->first_free[CROARING_ART_LEAF_TYPE];
+         i < art->capacities[CROARING_ART_LEAF_TYPE];
+         i = art_node_get_next_free(art,
+                                    art_to_ref(i, CROARING_ART_LEAF_TYPE))) {
+        num_free++;
+    }
+    return art->capacities[CROARING_ART_LEAF_TYPE] - num_free;
+}
+
 void art_printf(const art_t *art) {
     if (art->root == CROARING_ART_NULL_REF) {
         return;

src/roaring64.c

@@ -173,6 +173,24 @@ static inline leaf_t replace_container(roaring64_bitmap_t *r, leaf_t *leaf,
     return *leaf;
 }
 
+/**
+ * Grows the array of container pointers (and the parallel typecode array) to
+ * `new_capacity` entries. Does nothing if `new_capacity` is not larger than
+ * the current capacity.
+ */
+static void grow_containers_to(roaring64_bitmap_t *r, uint64_t new_capacity) {
+    if (new_capacity <= r->capacity) {
+        return;
+    }
+    uint64_t increase = new_capacity - r->capacity;
+    r->containers = (container_t **)roaring_realloc(
+        r->containers, new_capacity * sizeof(container_t *));
+    memset(r->containers + r->capacity, 0, increase * sizeof(container_t *));
+    r->typecodes = (uint8_t *)roaring_realloc(r->typecodes,
+                                              new_capacity * sizeof(uint8_t));
+    r->capacity = new_capacity;
+}
+
 /**
  * Extends the array of container pointers (and the parallel typecode array).
  */
@@ -189,13 +207,17 @@ static void extend_containers(roaring64_bitmap_t *r) {
     } else {
         new_capacity = 5 * r->capacity / 4;
     }
-    uint64_t increase = new_capacity - r->capacity;
-    r->containers = (container_t **)roaring_realloc(
-        r->containers, new_capacity * sizeof(container_t *));
-    memset(r->containers + r->capacity, 0, increase * sizeof(container_t *));
-    r->typecodes = (uint8_t *)roaring_realloc(r->typecodes,
-                                              new_capacity * sizeof(uint8_t));
-    r->capacity = new_capacity;
+    grow_containers_to(r, new_capacity);
+}
+
+/**
+ * Reserves room for `num_containers` containers in the ART and in the
+ * container array, so that a result of known maximum size can be built
+ * without repeatedly reallocating both. Intended for freshly created bitmaps.
+ */
+static void reserve_containers(roaring64_bitmap_t *r, uint64_t num_containers) {
+    art_reserve(&r->art, num_containers);
+    grow_containers_to(r, num_containers);
 }
 
 static uint64_t next_free_container_idx(const roaring64_bitmap_t *r) {
@@ -390,6 +412,7 @@ void roaring64_bitmap_free(roaring64_bitmap_t *r) {
 
 roaring64_bitmap_t *roaring64_bitmap_copy(const roaring64_bitmap_t *r) {
     roaring64_bitmap_t *result = roaring64_bitmap_create();
+    reserve_containers(result, art_num_leaves(&r->art));
 
     art_iterator_t it = art_init_iterator((art_t *)&r->art, /*first=*/true);
     while (it.value != NULL) {
@@ -1541,6 +1564,10 @@ double roaring64_bitmap_jaccard_index(const roaring64_bitmap_t *r1,
 roaring64_bitmap_t *roaring64_bitmap_or(const roaring64_bitmap_t *r1,
                                         const roaring64_bitmap_t *r2) {
     roaring64_bitmap_t *result = roaring64_bitmap_create();
+    // The result has at least max(n1, n2) containers and at most n1 + n2, so
+    // reserving the upper bound is at most a 2x overestimate.
+    reserve_containers(result,
+                       art_num_leaves(&r1->art) + art_num_leaves(&r2->art));
 
     art_iterator_t it1 = art_init_iterator((art_t *)&r1->art, /*first=*/true);
     art_iterator_t it2 = art_init_iterator((art_t *)&r2->art, /*first=*/true);

tests/art_unit.cpp

@@ -655,6 +655,68 @@ DEFINE_TEST(test_art_shrink_grow_node48) {
     art_free(&art);
 }
 
+DEFINE_TEST(test_art_num_leaves_and_reserve) {
+    std::vector<std::array<uint8_t, 6>> keys;
+    for (size_t i = 0; i < 300; ++i) {
+        keys.push_back({0, 0, 0, 0, static_cast<uint8_t>(i / 256),
+                        static_cast<uint8_t>(i % 256)});
+    }
+    {
+        // Counting through inserts and erasures, with a non-trivial free list.
+        art_t art;
+        art_init_cleared(&art);
+        assert_int_equal(art_num_leaves(&art), 0);
+        for (size_t i = 0; i < keys.size(); ++i) {
+            art_insert(&art, (art_key_chunk_t*)keys[i].data(), i);
+            assert_int_equal(art_num_leaves(&art), i + 1);
+        }
+        // Erase every third key, in an order that scatters the free list.
+        size_t erased = 0;
+        for (size_t i = 0; i < keys.size(); i += 3) {
+            art_val_t erased_val;
+            assert_true(
+                art_erase(&art, (art_key_chunk_t*)keys[i].data(), &erased_val));
+            ++erased;
+            assert_int_equal(art_num_leaves(&art), keys.size() - erased);
+        }
+        // Re-insert some, reusing free slots.
+        for (size_t i = 0; i < keys.size(); i += 6) {
+            art_insert(&art, (art_key_chunk_t*)keys[i].data(), i);
+            --erased;
+            assert_int_equal(art_num_leaves(&art), keys.size() - erased);
+        }
+        assert_art_valid(&art);
+        art_shrink_to_fit(&art);
+        assert_int_equal(art_num_leaves(&art), keys.size() - erased);
+        art_free(&art);
+    }
+    {
+        // Reserving up front leaves the count unchanged and the ART usable.
+        art_t art;
+        art_init_cleared(&art);
+        art_reserve(&art, keys.size());
+        assert_int_equal(art_num_leaves(&art), 0);
+        // Value pointers stay valid only if the leaf array is never
+        // reallocated, which is what the reservation guarantees.
+        art_val_t* first =
+            art_insert(&art, (art_key_chunk_t*)keys[0].data(), 0);
+        for (size_t i = 1; i < keys.size(); ++i) {
+            art_insert(&art, (art_key_chunk_t*)keys[i].data(), i);
+        }
+        assert_true(art_find(&art, (art_key_chunk_t*)keys[0].data()) == first);
+        // Reserving less than the current capacity is a no-op.
+        art_reserve(&art, 1);
+        assert_true(art_find(&art, (art_key_chunk_t*)keys[0].data()) == first);
+        assert_int_equal(art_num_leaves(&art), keys.size());
+        assert_art_valid(&art);
+        for (size_t i = 0; i < keys.size(); ++i) {
+            art_val_t* val = art_find(&art, (art_key_chunk_t*)keys[i].data());
+            assert_true(val != NULL && *val == i);
+        }
+        art_free(&art);
+    }
+}
+
 DEFINE_TEST(test_art_frozen_view) {
     {
         // ART with multiple node sizes.
@@ -747,6 +809,7 @@ int main() {
         cmocka_unit_test(test_art_iterator_insert),
         cmocka_unit_test(test_art_shadowed),
         cmocka_unit_test(test_art_shrink_grow_node48),
+        cmocka_unit_test(test_art_num_leaves_and_reserve),
         cmocka_unit_test(test_art_frozen_view),
     };
     return cmocka_run_group_tests(tests, NULL, NULL);