Add roaring_bitmap_word_range_to_bitset

include/roaring/roaring.h

@@ -588,6 +588,24 @@ void roaring_bitmap_to_uint32_array(const roaring_bitmap_t *r, uint32_t *ans);
  */
 bool roaring_bitmap_to_bitset(const roaring_bitmap_t *r, bitset_t *bitset);
 
+/**
+ * Writes the values of the bitmap in a range of 64-bit words, from
+ * 64 * first_word (included) to 64 * (first_word + n_words) (excluded), as a
+ * bitset in `words`: value 64 * (first_word + k) + i, if present, sets bit i
+ * of words[k]. All n_words words are written (the caller does not need to zero
+ * them). The cost is proportional to n_words plus the number of values and
+ * runs in the range, not to the size of the bitmap.
+ *
+ * first_word + n_words must be at most 2^26 (that is, 2^32 / 64).
+ *
+ *   uint64_t words[16];
+ *   // values 1024 to 2047, value 1024 + 64 * k + i in bit i of words[k]
+ *   roaring_bitmap_word_range_to_bitset(mybitmap, 16, 16, words);
+ */
+void roaring_bitmap_word_range_to_bitset(const roaring_bitmap_t *r,
+                                         uint32_t first_word, uint32_t n_words,
+                                         uint64_t *words);
+
 /**
  * Convert the bitmap to a sorted array from `offset` by `limit`, output in
  * `ans`.

src/roaring.c

@@ -3705,6 +3705,74 @@ roaring_bitmap_t *roaring_bitmap_portable_deserialize_frozen(const char *buf) {
 #endif
 }
 
+void roaring_bitmap_word_range_to_bitset(const roaring_bitmap_t *r,
+                                         uint32_t first_word, uint32_t n_words,
+                                         uint64_t *words) {
+    memset(words, 0, (size_t)n_words * sizeof(uint64_t));
+    if (n_words == 0) return;
+    assert((uint64_t)first_word + n_words <= (UINT64_C(1) << 26));
+    const uint64_t begin = (uint64_t)first_word * 64;
+    const uint64_t end = begin + (uint64_t)n_words * 64;  // excluded
+    const roaring_array_t *ra = &r->high_low_container;
+    int32_t i = ra_get_index(ra, (uint16_t)(begin >> 16));
+    if (i < 0) i = -i - 1;
+    for (; i < ra->size; ++i) {
+        const uint64_t chunk = (uint64_t)ra->keys[i] << 16;
+        if (chunk >= end) break;
+        // The window covers the words [lo, hi) of the container, which go
+        // to out[0, hi - lo).
+        const uint32_t lo =
+            chunk < begin ? (uint32_t)((begin - chunk) / 64) : 0;
+        const uint32_t hi =
+            chunk + 65536 > end ? (uint32_t)((end - chunk) / 64) : 1024;
+        uint64_t *out = words + (chunk / 64 + lo - first_word);
+        const uint32_t vlo = lo * 64, vhi = hi * 64;  // values [vlo, vhi)
+        uint8_t type = ra->typecodes[i];
+        const container_t *c =
+            container_unwrap_shared(ra->containers[i], &type);
+        switch (type) {
+            case BITSET_CONTAINER_TYPE:
+                memcpy(out, const_CAST_bitset(c)->words + lo,
+                       (size_t)(hi - lo) * sizeof(uint64_t));
+                break;
+            case ARRAY_CONTAINER_TYPE: {
+                const array_container_t *src = const_CAST_array(c);
+                int32_t k = 0;
+                if (vlo > 0) {
+                    k = binarySearch(src->array, src->cardinality,
+                                     (uint16_t)vlo);
+                    if (k < 0) k = -k - 1;
+                }
+                for (; k < src->cardinality && src->array[k] < vhi; ++k) {
+                    const uint32_t v = src->array[k] - vlo;
+                    out[v / 64] |= UINT64_C(1) << (v % 64);
+                }
+            } break;
+            case RUN_CONTAINER_TYPE: {
+                const run_container_t *src = const_CAST_run(c);
+                int32_t k = 0;
+                if (vlo > 0) {
+                    // The last run starting at or before vlo may reach it.
+                    k = interleavedBinarySearch(src->runs, src->n_runs,
+                                                (uint16_t)vlo);
+                    if (k < 0) k = -k - 2;
+                    if (k < 0) k = 0;
+                }
+                for (; k < src->n_runs && src->runs[k].value < vhi; ++k) {
+                    uint32_t s = src->runs[k].value;
+                    uint32_t e = s + src->runs[k].length + 1;  // excluded
+                    if (e <= vlo) continue;
+                    if (s < vlo) s = vlo;
+                    if (e > vhi) e = vhi;
+                    bitset_set_range(out, s - vlo, e - vlo);
+                }
+            } break;
+            default:
+                roaring_unreachable;
+        }
+    }
+}
+
 bool roaring_bitmap_to_bitset(const roaring_bitmap_t *r, bitset_t *bitset) {
     uint32_t max_value = roaring_bitmap_maximum(r);
     size_t new_array_size = (size_t)(max_value / 64 + 1);

tests/toplevel_unit.c

@@ -3621,6 +3621,97 @@ DEFINE_TEST(test_flip_sparse_array_container) {
     }
 }
 
+// Checks roaring_bitmap_word_range_to_bitset against roaring_bitmap_contains.
+static void check_word_range_to_bitset(const roaring_bitmap_t *bm,
+                                       uint32_t first_word, uint32_t n_words) {
+    uint64_t *words = (uint64_t *)malloc((n_words + 2) * sizeof(uint64_t));
+    assert_non_null(words);
+    // Guard words around the output must not be touched.
+    words[0] = words[n_words + 1] = UINT64_C(0x0123456789abcdef);
+    for (uint32_t k = 1; k <= n_words; k++) words[k] = ~UINT64_C(0);
+    roaring_bitmap_word_range_to_bitset(bm, first_word, n_words, words + 1);
+    assert_true(words[0] == UINT64_C(0x0123456789abcdef));
+    assert_true(words[n_words + 1] == UINT64_C(0x0123456789abcdef));
+    for (uint32_t k = 0; k < n_words; k++) {
+        uint64_t expected = 0;
+        for (uint32_t i = 0; i < 64; i++) {
+            uint32_t v = (uint32_t)(((uint64_t)first_word + k) * 64 + i);
+            if (roaring_bitmap_contains(bm, v)) expected |= UINT64_C(1) << i;
+        }
+        assert_true(words[k + 1] == expected);
+    }
+    free(words);
+}
+
+DEFINE_TEST(test_word_range_to_bitset) {
+    roaring_bitmap_t *bm = roaring_bitmap_create();
+    uint32_t seed = 1234;
+    // Chunk 0: array, chunk 1: bitset, chunk 2: runs, chunk 4: array, and
+    // the last chunk of the 32-bit range: runs and values at the very end.
+    for (uint32_t k = 0; k < 1000; k++) {
+        seed = seed * 1664525 + 1013904223;
+        roaring_bitmap_add(bm, (seed >> 16) % 65536);
+    }
+    for (uint32_t k = 0; k < 30000; k++) {
+        seed = seed * 1664525 + 1013904223;
+        roaring_bitmap_add(bm, 65536 + (seed >> 16) % 65536);
+    }
+    for (uint32_t k = 0; k < 300; k++) {
+        seed = seed * 1664525 + 1013904223;
+        uint32_t s = 2 * 65536 + (seed >> 16) % 65000;
+        roaring_bitmap_add_range(bm, s, s + 1 + (seed >> 8) % 300);
+    }
+    roaring_bitmap_add(bm, 4 * 65536 + 63);
+    roaring_bitmap_add(bm, 4 * 65536 + 64);
+    roaring_bitmap_add_range(bm, UINT64_C(0xFFFF0000) + 100,
+                             UINT64_C(0xFFFF0000) + 40000);
+    roaring_bitmap_add_range(bm, UINT64_C(0xFFFFFFF0), UINT64_C(0x100000000));
+    roaring_bitmap_run_optimize(bm);
+    roaring_statistics_t stats;
+    roaring_bitmap_statistics(bm, &stats);
+    assert_true(stats.n_array_containers >= 2);
+    assert_true(stats.n_bitset_containers >= 1);
+    assert_true(stats.n_run_containers >= 2);
+
+    const uint32_t last = (UINT32_C(1) << 26);
+    const uint32_t windows[][2] = {{0, 0},
+                                   {0, 1},
+                                   {0, 1024},
+                                   {0, 5000},
+                                   {1, 1},
+                                   {3, 2},
+                                   {1023, 2},
+                                   {1000, 100},
+                                   {1024, 1024},
+                                   {1500, 1},
+                                   {2047, 3},
+                                   {2048, 1024},
+                                   {2100, 7},
+                                   {2900, 300},
+                                   {4096, 1},
+                                   {4096, 2},
+                                   {3000, 2000},
+                                   {10000, 10},
+                                   {last - 1, 1},
+                                   {last - 1024, 1024},
+                                   {last - 700, 700},
+                                   {last - 1100, 200},
+                                   {last - 3000, 3000}};
+    for (size_t w = 0; w < sizeof(windows) / sizeof(windows[0]); w++) {
+        check_word_range_to_bitset(bm, windows[w][0], windows[w][1]);
+    }
+    for (int trial = 0; trial < 200; trial++) {
+        seed = seed * 1664525 + 1013904223;
+        uint32_t first = (seed >> 8) % 5200;
+        seed = seed * 1664525 + 1013904223;
+        check_word_range_to_bitset(bm, first, (seed >> 8) % 1500);
+    }
+    roaring_bitmap_t *empty = roaring_bitmap_create();
+    check_word_range_to_bitset(empty, 100, 50);
+    roaring_bitmap_free(empty);
+    roaring_bitmap_free(bm);
+}
+
 // randomized test for rank query
 DEFINE_TEST(select_test) {
     const int min_runs = 1;
@@ -6009,6 +6100,7 @@ int main() {
         cmocka_unit_test(test_flip_run_container_removal),
         cmocka_unit_test(test_flip_run_container_removal2),
         cmocka_unit_test(test_flip_sparse_array_container),
+        cmocka_unit_test(test_word_range_to_bitset),
         cmocka_unit_test(select_test),
         cmocka_unit_test(test_subset),
         cmocka_unit_test(test_or_many_memory_leak),