Add some missing features to 64-bit bitmaps (#146)

* Add shrink_to_fit() to BitMap64
* Add shift() to BitMap64
* Add overwrite() to BitMap64

Author: Ezibenroc

pyroaring/abstract_bitmap.pxi

@@ -856,6 +856,7 @@ cdef class AbstractBitMap64:
             self._h_val = 0
         if optimize:
             self.run_optimize()
+            self.shrink_to_fit()
 
     def __init__(self, values=None, copy_on_write=False, optimize=True):
         """
@@ -895,6 +896,9 @@ cdef class AbstractBitMap64:
     def run_optimize(self):
         return croaring.roaring64_bitmap_run_optimize(self._c_bitmap)
 
+    def shrink_to_fit(self):
+        return croaring.roaring64_bitmap_shrink_to_fit(self._c_bitmap)
+
     def __dealloc__(self):
         if self._c_bitmap is not NULL:
             croaring.roaring64_bitmap_free(self._c_bitmap)
@@ -1044,6 +1048,21 @@ cdef class AbstractBitMap64:
         """
         return self.from_ptr(croaring.roaring64_bitmap_flip(self._c_bitmap, start, end))
 
+    def shift(self, int64_t offset):
+        """
+        Add the value 'offset' to each and every value of the bitmap.
+
+        If offset + element is outside of the range [0,2^64), that the element will be dropped.
+
+        >>> bm = BitMap64([3, 12])
+        >>> bm.shift(21)
+        BitMap64([24, 33])
+        """
+        if offset >= 0:
+            return self.from_ptr(croaring.roaring64_bitmap_add_offset_signed(self._c_bitmap, True, <uint64_t>offset))
+        else:
+            return self.from_ptr(croaring.roaring64_bitmap_add_offset_signed(self._c_bitmap, False, <uint64_t>0 - <uint64_t>offset))
+
     def get_statistics(self):
         """
         Return relevant metrics about the bitmap.

pyroaring/bitmap.pxi

@@ -441,6 +441,23 @@ cdef class BitMap64(AbstractBitMap64):
         """
         self.__ixor__(other)
 
+    def overwrite(self, AbstractBitMap64 other):
+        """
+        Clear the bitmap and overwrite it with another.
+
+        >>> bm = BitMap64([3, 12])
+        >>> other = BitMap64([4, 14])
+        >>> bm.overwrite(other)
+        >>> other.remove(4)
+        >>> bm
+        BitMap64([4, 14])
+        >>> other
+        BitMap64([14])
+        """
+        if self._c_bitmap == other._c_bitmap:
+            raise ValueError('Cannot overwrite itself')
+        croaring.roaring64_bitmap_overwrite(self._c_bitmap, other._c_bitmap)
+
     def clear(self):
         """
         Remove all elements from this set.

pyroaring/croaring.pxd

@@ -142,6 +142,7 @@ cdef extern from "roaring.h":
     uint64_t roaring64_bitmap_minimum(const roaring64_bitmap_t *r)
     uint64_t roaring64_bitmap_maximum(const roaring64_bitmap_t *r)
     bool roaring64_bitmap_run_optimize(roaring64_bitmap_t *r)
+    size_t roaring64_bitmap_shrink_to_fit(roaring64_bitmap_t *r)
     size_t roaring64_bitmap_size_in_bytes(const roaring64_bitmap_t *r)
     bool roaring64_bitmap_equals(const roaring64_bitmap_t *r1, const roaring64_bitmap_t *r2)
     bool roaring64_bitmap_is_subset(const roaring64_bitmap_t *r1, const roaring64_bitmap_t *r2)
@@ -166,6 +167,8 @@ cdef extern from "roaring.h":
     size_t roaring64_bitmap_portable_deserialize_size(const char *buf, size_t maxbytes)
     roaring64_bitmap_t *roaring64_bitmap_portable_deserialize_safe(const char *buf, size_t maxbytes)
     bool roaring64_bitmap_internal_validate(const roaring64_bitmap_t *r, const char **reason)
+    roaring64_bitmap_t *roaring64_bitmap_add_offset_signed(const roaring64_bitmap_t *r, bool positive, uint64_t offset)
+    void roaring64_bitmap_overwrite(roaring64_bitmap_t *dest, const roaring64_bitmap_t *src)
     roaring64_iterator_t *roaring64_iterator_create(const roaring64_bitmap_t *r)
     void roaring64_iterator_free(roaring64_iterator_t *it)
     bool roaring64_iterator_has_value(const roaring64_iterator_t *it)

test.py

@@ -1090,7 +1090,6 @@ def test_flip_inplace(
         bm_after.flip_inplace(start, end)
         self.check_flip(bm_before, bm_after, start, end)
 
-@pytest.mark.skipif(not is_32_bits, reason="not supported yet")
 class TestShift(Util):
     @given(bitmap_cls, hyp_collection, int64, st.booleans())
     def test_shift(
@@ -1104,7 +1103,8 @@ def test_shift(
         bm_copy = cls(bm_before)
         bm_after = bm_before.shift(offset)
         assert bm_before == bm_copy
-        expected = cls([val + offset for val in values if val + offset in range(0, 2**32)], copy_on_write=cow)
+        upper = 2**32 if is_32_bits else 2**64
+        expected = cls([val + offset for val in values if val + offset in range(0, upper)], copy_on_write=cow)
         assert bm_after == expected
 
 @pytest.mark.skipif(not is_32_bits, reason="not supported yet")
@@ -1297,7 +1297,6 @@ def test_run_optimize(self, cls: type[EitherBitMap]) -> None:
         bm3 = cls(bm1)  # optimize is True by default
         assert stats == bm3.get_statistics()
 
-    @pytest.mark.skipif(not is_32_bits, reason="not supported yet")
     @given(bitmap_cls)
     def test_shrink_to_fit(self, cls: type[EitherBitMap]) -> None:
         bm1 = BitMap()
@@ -1561,7 +1560,6 @@ def test_copy_writable(self, list1: list[int], cow: bool) -> None:
         assert new_element in b2
         assert new_element not in b1
 
-    @pytest.mark.skipif(not is_32_bits, reason="not supported yet")
     @given(bitmap_cls, small_integer_list, small_integer_list, st.booleans())
     def test_overwrite(self, BitMapClass: type[EitherBitMap], list1: list[int], list2: list[int], cow: bool) -> None:
         assume(set(list1) != set(list2))