From cc63c17f06ef86a08aa605c7876388458f5c1545 Mon Sep 17 00:00:00 2001 From: Daniel Lemire Date: Mon, 21 Sep 2026 19:04:08 -0400 Subject: [PATCH] minor optimizations to roaring64_bitmap_or --- include/roaring/art/art.h | 15 ++++++++++ src/art/art.c | 48 +++++++++++++++++++++++------ src/roaring64.c | 41 ++++++++++++++++++++----- tests/art_unit.cpp | 63 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 16 deletions(-) diff --git a/include/roaring/art/art.h b/include/roaring/art/art.h index 16b1e5516..ef13b339c 100644 --- a/include/roaring/art/art.h +++ b/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. diff --git a/src/art/art.c b/src/art/art.c index 5235be80b..6140b1e0a 100644 --- a/src/art/art.c +++ b/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; diff --git a/src/roaring64.c b/src/roaring64.c index f303ae6e3..823840906 100644 --- a/src/roaring64.c +++ b/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); diff --git a/tests/art_unit.cpp b/tests/art_unit.cpp index fd77e2046..84a83c88b 100644 --- a/tests/art_unit.cpp +++ b/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> keys; + for (size_t i = 0; i < 300; ++i) { + keys.push_back({0, 0, 0, 0, static_cast(i / 256), + static_cast(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);