From 6467be6a132350cfc0c55c94aec41bfc4e8ae62c Mon Sep 17 00:00:00 2001
From: Zachary Dremann <dremann@gmail.com>
Date: Sun, 4 Aug 2024 14:56:11 -0400
Subject: [PATCH] Add `roaring64_bitmap_steal_roaring32` function

This allows creating a 64 bitmap cheaply from a roaring_bitmap_t. Ideally, we
should add a `roaring64_bitmap_add_offset` too.
---
 include/roaring/roaring64.h |  9 ++++++
 src/roaring64.c             | 55 ++++++++++++++++++++++++++-----------
 2 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/include/roaring/roaring64.h b/include/roaring/roaring64.h
index c71bccb91..efc828352 100644
--- a/include/roaring/roaring64.h
+++ b/include/roaring/roaring64.h
@@ -1,6 +1,7 @@
 #ifndef ROARING64_H
 #define ROARING64_H
 
+#include <roaring.h>
 #include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
@@ -92,6 +93,14 @@ roaring64_bitmap_t *roaring64_bitmap_of_ptr(size_t n_args,
         &((const uint64_t[]){0, __VA_ARGS__})[1])
 #endif
 
+/**
+ * Create a new bitmap by stealing containers from a 32 bit roaring bitmap.
+ *
+ * After calling this function, the original bitmap will be empty, and the
+ * returned bitmap will contain all the values from the original bitmap.
+ */
+roaring64_bitmap_t *roaring64_bitmap_steal_roaring32(roaring_bitmap_t *r);
+
 /**
  * Create a new bitmap containing all the values in [min, max) that are at a
  * distance k*step from min.
diff --git a/src/roaring64.c b/src/roaring64.c
index d41507b3c..2ab13efb0 100644
--- a/src/roaring64.c
+++ b/src/roaring64.c
@@ -178,6 +178,43 @@ roaring64_bitmap_t *roaring64_bitmap_copy(const roaring64_bitmap_t *r) {
     return result;
 }
 
+/**
+ * Steal the containers from a 32-bit bitmap and insert them into a 64-bit
+ * bitmap (with an offset)
+ *
+ * After calling this function, the original bitmap will be empty, and the
+ * returned bitmap will contain all the values from the original bitmap.
+ */
+static void steal_from_roaring32_offset(roaring64_bitmap_t *dst,
+                                        roaring_bitmap_t *src,
+                                        uint32_t high_bits) {
+    uint64_t key_base = ((uint64_t)high_bits) << 32;
+    uint32_t r32_size = ra_get_size(&src->high_low_container);
+    for (uint32_t i = 0; i < r32_size; ++i) {
+        uint16_t key = ra_get_key_at_index(&src->high_low_container, i);
+        uint8_t typecode;
+        container_t *container = ra_get_container_at_index(
+            &src->high_low_container, (uint16_t)i, &typecode);
+
+        uint8_t high48[ART_KEY_BYTES];
+        uint64_t high48_bits = key_base | ((uint64_t)key << 16);
+        split_key(high48_bits, high48);
+        leaf_t *leaf = create_leaf(container, typecode);
+        art_insert(&dst->art, high48, (art_val_t *)leaf);
+    }
+    // We stole all the containers, so leave behind a size of zero
+    src->high_low_container.size = 0;
+}
+
+roaring64_bitmap_t *roaring64_bitmap_steal_roaring32(
+    roaring_bitmap_t *bitmap32) {
+    roaring64_bitmap_t *result = roaring64_bitmap_create();
+
+    steal_from_roaring32_offset(result, bitmap32, 0);
+
+    return result;
+}
+
 roaring64_bitmap_t *roaring64_bitmap_from_range(uint64_t min, uint64_t max,
                                                 uint64_t step) {
     if (step == 0 || max <= min) {
@@ -1947,22 +1984,8 @@ roaring64_bitmap_t *roaring64_bitmap_portable_deserialize_safe(
         read_bytes += bitmap32_size;
 
         // Insert all containers of the 32-bit bitmap into the 64-bit bitmap.
-        uint32_t r32_size = ra_get_size(&bitmap32->high_low_container);
-        for (size_t i = 0; i < r32_size; ++i) {
-            uint16_t key16 =
-                ra_get_key_at_index(&bitmap32->high_low_container, (uint16_t)i);
-            uint8_t typecode;
-            container_t *container = ra_get_container_at_index(
-                &bitmap32->high_low_container, (uint16_t)i, &typecode);
-
-            uint64_t high48_bits =
-                (((uint64_t)high32) << 32) | (((uint64_t)key16) << 16);
-            uint8_t high48[ART_KEY_BYTES];
-            split_key(high48_bits, high48);
-            leaf_t *leaf = create_leaf(container, typecode);
-            art_insert(&r->art, high48, (art_val_t *)leaf);
-        }
-        roaring_bitmap_free_without_containers(bitmap32);
+        steal_from_roaring32_offset(r, bitmap32, high32);
+        roaring_bitmap_free(bitmap32);
     }
     return r;
 }