Add roaring64_iterator_{reinit, reinit_last}

These allow reusing an existing iterator.
RoaringBitmap · Jan 16, 2024 · f18a415 · f18a415
1 parent dfb10b7
commit f18a415
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 6 deletions.
diff --git a/include/roaring/roaring64.h b/include/roaring/roaring64.h
@@ -402,6 +402,20 @@ roaring64_iterator_t *roaring64_iterator_create(const roaring64_bitmap_t *r);
 roaring64_iterator_t *roaring64_iterator_create_last(
     const roaring64_bitmap_t *r);
 
+/**
+ * Re-initializes an existing iterator. Functionally the same as
+ * `roaring64_iterator_create` without a allocation.
+ */
+void roaring64_iterator_reinit(const roaring64_bitmap_t *r,
+                               roaring64_iterator_t *it);
+
+/**
+ * Re-initializes an existing iterator. Functionally the same as
+ * `roaring64_iterator_create_last` without a allocation.
+ */
+void roaring64_iterator_reinit_last(const roaring64_bitmap_t *r,
+                                    roaring64_iterator_t *it);
+
 /**
  * Creates a copy of the iterator. Caller is responsible for calling
  * `roaring64_iterator_free()` on the resulting iterator.

diff --git a/src/roaring64.c b/src/roaring64.c
@@ -1361,10 +1361,8 @@ static inline bool roaring64_iterator_init_at_leaf_last(
     return (it->has_value = true);
 }
 
-static inline roaring64_iterator_t *roaring64_iterator_create_at(
-    const roaring64_bitmap_t *r, bool first) {
-    roaring64_iterator_t *it =
-        (roaring64_iterator_t *)roaring_malloc(sizeof(roaring64_iterator_t));
+static inline roaring64_iterator_t *roaring64_iterator_init_at(
+    const roaring64_bitmap_t *r, roaring64_iterator_t *it, bool first) {
     it->parent = r;
     it->art_it = art_init_iterator(&r->art, first);
     it->has_value = it->art_it.value != NULL;
@@ -1379,12 +1377,26 @@ static inline roaring64_iterator_t *roaring64_iterator_create_at(
 }
 
 roaring64_iterator_t *roaring64_iterator_create(const roaring64_bitmap_t *r) {
-    return roaring64_iterator_create_at(r, /*first=*/true);
+    roaring64_iterator_t *it =
+        (roaring64_iterator_t *)roaring_malloc(sizeof(roaring64_iterator_t));
+    return roaring64_iterator_init_at(r, it, /*first=*/true);
 }
 
 roaring64_iterator_t *roaring64_iterator_create_last(
     const roaring64_bitmap_t *r) {
-    return roaring64_iterator_create_at(r, /*first=*/false);
+    roaring64_iterator_t *it =
+        (roaring64_iterator_t *)roaring_malloc(sizeof(roaring64_iterator_t));
+    return roaring64_iterator_init_at(r, it, /*first=*/false);
+}
+
+void roaring64_iterator_reinit(const roaring64_bitmap_t *r,
+                               roaring64_iterator_t *it) {
+    roaring64_iterator_init_at(r, it, /*first=*/true);
+}
+
+void roaring64_iterator_reinit_last(const roaring64_bitmap_t *r,
+                                    roaring64_iterator_t *it) {
+    roaring64_iterator_init_at(r, it, /*first=*/false);
 }
 
 roaring64_iterator_t *roaring64_iterator_copy(const roaring64_iterator_t *it) {

diff --git a/tests/roaring64_unit.cpp b/tests/roaring64_unit.cpp
@@ -937,6 +937,56 @@ DEFINE_TEST(test_iterator_create_last) {
     roaring64_bitmap_free(r);
 }
 
+DEFINE_TEST(test_iterator_reinit) {
+    roaring64_bitmap_t* r = roaring64_bitmap_create();
+
+    roaring64_bitmap_add(r, 0);
+    roaring64_bitmap_add(r, 1ULL << 35);
+    roaring64_bitmap_add(r, (1ULL << 35) + 1);
+    roaring64_bitmap_add(r, (1ULL << 35) + 2);
+    roaring64_bitmap_add(r, (1ULL << 36));
+
+    roaring64_iterator_t* it = roaring64_iterator_create(r);
+    assert_true(roaring64_iterator_advance(it));
+    assert_true(roaring64_iterator_advance(it));
+    assert_true(roaring64_iterator_advance(it));
+    assert_true(roaring64_iterator_previous(it));
+    assert_true(roaring64_iterator_has_value(it));
+    assert_int_equal(roaring64_iterator_value(it), ((1ULL << 35) + 1));
+
+    roaring64_iterator_reinit(r, it);
+    assert_true(roaring64_iterator_has_value(it));
+    assert_int_equal(roaring64_iterator_value(it), 0);
+
+    roaring64_iterator_free(it);
+    roaring64_bitmap_free(r);
+}
+
+DEFINE_TEST(test_iterator_reinit_last) {
+    roaring64_bitmap_t* r = roaring64_bitmap_create();
+
+    roaring64_bitmap_add(r, 0);
+    roaring64_bitmap_add(r, 1ULL << 35);
+    roaring64_bitmap_add(r, (1ULL << 35) + 1);
+    roaring64_bitmap_add(r, (1ULL << 35) + 2);
+    roaring64_bitmap_add(r, (1ULL << 36));
+
+    roaring64_iterator_t* it = roaring64_iterator_create(r);
+    assert_true(roaring64_iterator_advance(it));
+    assert_true(roaring64_iterator_advance(it));
+    assert_true(roaring64_iterator_advance(it));
+    assert_true(roaring64_iterator_previous(it));
+    assert_true(roaring64_iterator_has_value(it));
+    assert_int_equal(roaring64_iterator_value(it), ((1ULL << 35) + 1));
+
+    roaring64_iterator_reinit_last(r, it);
+    assert_true(roaring64_iterator_has_value(it));
+    assert_int_equal(roaring64_iterator_value(it), (1ULL << 36));
+
+    roaring64_iterator_free(it);
+    roaring64_bitmap_free(r);
+}
+
 DEFINE_TEST(test_iterator_copy) {
     roaring64_bitmap_t* r = roaring64_bitmap_create();
 
@@ -1138,6 +1188,8 @@ int main() {
         cmocka_unit_test(test_iterate),
         cmocka_unit_test(test_iterator_create),
         cmocka_unit_test(test_iterator_create_last),
+        cmocka_unit_test(test_iterator_reinit),
+        cmocka_unit_test(test_iterator_reinit_last),
         cmocka_unit_test(test_iterator_copy),
         cmocka_unit_test(test_iterator_advance),
         cmocka_unit_test(test_iterator_previous),