Introduce Iter::next_range and next_range_back

fuzz/fuzz_targets/arbitrary_ops/mod.rs

@@ -103,6 +103,8 @@ pub enum BitmapIteratorOperation {
     AdvanceBackTo(Num),
     Nth(Num),
     NthBack(Num),
+    NextRange,
+    NextRangeBack,
 }
 
 impl ReadBitmapOperation {
@@ -466,6 +468,69 @@ impl<'a> CRoaringIterRange<'a> {
         }
         self.next_back()
     }
+
+    fn next_range(&mut self) -> Option<RangeInclusive<u32>> {
+        if self.empty {
+            return None;
+        }
+        self.cursor.reset_at_or_after(self.start);
+        let range_start = self.cursor.current()?;
+        if range_start > self.end_inclusive {
+            self.empty = true;
+            return None;
+        }
+        let mut range_end_inclusive = range_start;
+        while range_end_inclusive < self.end_inclusive {
+            if let Some(next) = self.cursor.next() {
+                if next == range_end_inclusive + 1 {
+                    range_end_inclusive = next;
+                    continue;
+                }
+            }
+            break;
+        }
+
+        if range_end_inclusive == self.end_inclusive {
+            self.empty = true;
+        } else {
+            self.start = range_end_inclusive + 1;
+        }
+
+        Some(range_start..=range_end_inclusive)
+    }
+
+    fn next_range_back(&mut self) -> Option<RangeInclusive<u32>> {
+        if self.empty {
+            return None;
+        }
+        self.cursor.reset_at_or_after(self.end_inclusive);
+        if self.cursor.current().is_none_or(|n| n > self.end_inclusive) {
+            self.cursor.move_prev();
+        }
+        let range_end_inclusive = self.cursor.current()?;
+        if range_end_inclusive < self.start {
+            self.empty = true;
+            return None;
+        }
+        let mut range_start = range_end_inclusive;
+        while range_start > self.start {
+            if let Some(prev) = self.cursor.prev() {
+                if prev == range_start - 1 {
+                    range_start = prev;
+                    continue;
+                }
+            }
+            break;
+        }
+
+        if range_start == self.start {
+            self.empty = true;
+        } else {
+            self.end_inclusive = range_start - 1;
+        }
+
+        Some(range_start..=range_end_inclusive)
+    }
 }
 
 impl BitmapIteratorOperation {
@@ -491,6 +556,12 @@ impl BitmapIteratorOperation {
             BitmapIteratorOperation::NthBack(n) => {
                 assert_eq!(x.nth_back(n.0), y.nth_back(n.0 as usize));
             }
+            BitmapIteratorOperation::NextRange => {
+                assert_eq!(x.next_range(), y.next_range());
+            }
+            BitmapIteratorOperation::NextRangeBack => {
+                assert_eq!(x.next_range_back(), y.next_range_back());
+            }
         }
     }
 }

roaring/src/bitmap/container.rs

@@ -438,13 +438,41 @@ impl DoubleEndedIterator for Iter<'_> {
 impl ExactSizeIterator for Iter<'_> {}
 
 impl Iter<'_> {
+    pub(crate) fn peek(&self) -> Option<u32> {
+        self.inner.peek().map(|i| util::join(self.key, i))
+    }
+
+    pub(crate) fn peek_back(&self) -> Option<u32> {
+        self.inner.peek_back().map(|i| util::join(self.key, i))
+    }
+
     pub(crate) fn advance_to(&mut self, index: u16) {
         self.inner.advance_to(index);
     }
 
     pub(crate) fn advance_back_to(&mut self, index: u16) {
         self.inner.advance_back_to(index);
     }
+
+    /// Returns the range of consecutive set bits from the current position to the end of the current run
+    ///
+    /// After this call, the iterator will be positioned at the first item after the returned range.
+    /// Returns `None` if the iterator is exhausted.
+    pub(crate) fn next_range(&mut self) -> Option<RangeInclusive<u32>> {
+        self.inner
+            .next_range()
+            .map(|r| util::join(self.key, *r.start())..=util::join(self.key, *r.end()))
+    }
+
+    /// Returns the range of consecutive set bits from the start of the current run to the current back position
+    ///
+    /// After this call, the back of the iterator will be positioned at the last item before the returned range.
+    /// Returns `None` if the iterator is exhausted.
+    pub(crate) fn next_range_back(&mut self) -> Option<RangeInclusive<u32>> {
+        self.inner
+            .next_range_back()
+            .map(|r| util::join(self.key, *r.start())..=util::join(self.key, *r.end()))
+    }
 }
 
 impl fmt::Debug for Container {

roaring/src/bitmap/iter.rs

@@ -92,6 +92,100 @@ fn advance_to_impl<'a, It>(
     }
 }
 
+fn next_range_impl<'a, It>(
+    front_iter: &mut Option<container::Iter<'a>>,
+    containers: &mut It,
+    back_iter: &mut Option<container::Iter<'a>>,
+) -> Option<core::ops::RangeInclusive<u32>>
+where
+    It: Iterator + Clone,
+    It: AsRef<[Container]>,
+    It::Item: IntoIterator<IntoIter = container::Iter<'a>>,
+{
+    let range = loop {
+        if let Some(r) = and_then_or_clear(front_iter, container::Iter::next_range) {
+            break r;
+        }
+        *front_iter = match containers.next() {
+            Some(inner) => Some(inner.into_iter()),
+            None => return and_then_or_clear(back_iter, container::Iter::next_range),
+        }
+    };
+    let (range_start, mut range_end) = (*range.start(), *range.end());
+    while range_end & 0xFFFF == 0xFFFF {
+        let Some(after_end) = range_end.checked_add(1) else {
+            return Some(range_start..=range_end);
+        };
+        let (next_key, _) = util::split(after_end);
+
+        if containers.as_ref().first().is_some_and(|c| c.key == next_key && c.contains(0)) {
+            let mut iter = containers.next().unwrap().into_iter();
+            let next_range = iter.next_range().unwrap();
+            *front_iter = Some(iter);
+            debug_assert_eq!(*next_range.start(), after_end);
+            range_end = *next_range.end();
+        } else {
+            if let Some(iter) = back_iter {
+                if iter.peek() == Some(after_end) {
+                    let next_range = iter.next_range().unwrap();
+                    debug_assert_eq!(*next_range.start(), after_end);
+                    range_end = *next_range.end();
+                }
+            }
+            break;
+        }
+    }
+
+    Some(range_start..=range_end)
+}
+
+fn next_range_back_impl<'a, It>(
+    front_iter: &mut Option<container::Iter<'a>>,
+    containers: &mut It,
+    back_iter: &mut Option<container::Iter<'a>>,
+) -> Option<core::ops::RangeInclusive<u32>>
+where
+    It: DoubleEndedIterator,
+    It: AsRef<[Container]>,
+    It::Item: IntoIterator<IntoIter = container::Iter<'a>>,
+{
+    let range = loop {
+        if let Some(r) = and_then_or_clear(back_iter, container::Iter::next_range_back) {
+            break r;
+        }
+        *back_iter = match containers.next_back() {
+            Some(inner) => Some(inner.into_iter()),
+            None => return and_then_or_clear(front_iter, container::Iter::next_range_back),
+        }
+    };
+    let (mut range_start, range_end) = (*range.start(), *range.end());
+    while range_start & 0xFFFF == 0 {
+        let Some(before_start) = range_start.checked_sub(1) else {
+            return Some(range_start..=range_end);
+        };
+        let (prev_key, _) = util::split(before_start);
+
+        if containers.as_ref().last().is_some_and(|c| c.key == prev_key && c.contains(u16::MAX)) {
+            let mut iter = containers.next_back().unwrap().into_iter();
+            let next_range = iter.next_range_back().unwrap();
+            *back_iter = Some(iter);
+            debug_assert_eq!(*next_range.end(), before_start);
+            range_start = *next_range.start();
+        } else {
+            if let Some(iter) = front_iter {
+                if iter.key == prev_key && iter.peek_back() == Some(before_start) {
+                    let next_range = iter.next_range_back().unwrap();
+                    debug_assert_eq!(*next_range.end(), before_start);
+                    range_start = *next_range.start();
+                }
+            }
+            break;
+        }
+    }
+
+    Some(range_start..=range_end)
+}
+
 fn advance_back_to_impl<'a, It>(
     n: u32,
     front_iter: &mut Option<container::Iter<'a>>,
@@ -197,6 +291,46 @@ impl Iter<'_> {
     pub fn advance_back_to(&mut self, n: u32) {
         advance_back_to_impl(n, &mut self.front, &mut self.containers, &mut self.back);
     }
+
+    /// Returns the range of consecutive set bits from the current position to the end of the current run
+    ///
+    /// After this call, the iterator will be positioned at the first item after the returned range.
+    /// Returns `None` if the iterator is exhausted.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use roaring::RoaringBitmap;
+    ///
+    /// let bm = RoaringBitmap::from([1, 2, 4, 5]);
+    /// let mut iter = bm.iter();
+    /// assert_eq!(iter.next_range(), Some(1..=2));
+    /// assert_eq!(iter.next(), Some(4));
+    /// assert_eq!(iter.next_range(), Some(5..=5));
+    /// ```
+    pub fn next_range(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
+        next_range_impl(&mut self.front, &mut self.containers, &mut self.back)
+    }
+
+    /// Returns the range of consecutive set bits from the start of the current run to the current back position
+    ///
+    /// After this call, the back of the iterator will be positioned at the last item before the returned range.
+    /// Returns `None` if the iterator is exhausted.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use roaring::RoaringBitmap;
+    ///
+    /// let bm = RoaringBitmap::from([1, 2, 4, 5]);
+    /// let mut iter = bm.iter();
+    /// assert_eq!(iter.next_range_back(), Some(4..=5));
+    /// assert_eq!(iter.next_back(), Some(2));
+    /// assert_eq!(iter.next_range_back(), Some(1..=1));
+    /// ```
+    pub fn next_range_back(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
+        next_range_back_impl(&mut self.front, &mut self.containers, &mut self.back)
+    }
 }
 
 impl IntoIter {
@@ -245,6 +379,46 @@ impl IntoIter {
     pub fn advance_back_to(&mut self, n: u32) {
         advance_back_to_impl(n, &mut self.front, &mut self.containers, &mut self.back);
     }
+
+    /// Returns the range of consecutive set bits from the current position to the end of the current run
+    ///
+    /// After this call, the iterator will be positioned at the first item after the returned range.
+    /// Returns `None` if the iterator is exhausted.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use roaring::RoaringBitmap;
+    ///
+    /// let bm = RoaringBitmap::from([1, 2, 4, 5]);
+    /// let mut iter = bm.into_iter();
+    /// assert_eq!(iter.next_range(), Some(1..=2));
+    /// assert_eq!(iter.next(), Some(4));
+    /// assert_eq!(iter.next_range(), Some(5..=5));
+    /// ```
+    pub fn next_range(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
+        next_range_impl(&mut self.front, &mut self.containers, &mut self.back)
+    }
+
+    /// Returns the range of consecutive set bits from the start of the current run to the current back position
+    ///
+    /// After this call, the back of the iterator will be positioned at the last item before the returned range.
+    /// Returns `None` if the iterator is exhausted.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use roaring::RoaringBitmap;
+    ///
+    /// let bm = RoaringBitmap::from([1, 2, 4, 5]);
+    /// let mut iter = bm.into_iter();
+    /// assert_eq!(iter.next_range_back(), Some(4..=5));
+    /// assert_eq!(iter.next_back(), Some(2));
+    /// assert_eq!(iter.next_range_back(), Some(1..=1));
+    /// ```
+    pub fn next_range_back(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
+        next_range_back_impl(&mut self.front, &mut self.containers, &mut self.back)
+    }
 }
 
 fn size_hint_impl(

roaring/src/bitmap/store/array_store/mod.rs

@@ -25,6 +25,38 @@ pub(crate) struct ArrayStore {
     vec: Vec<u16>,
 }
 
+/// Return the first contiguous range of elements in a sorted slice.
+pub(crate) fn first_contiguous_range_len(slice: &[u16]) -> usize {
+    let [first, rest @ ..] = slice else {
+        // Explicitly empty range
+        return 0;
+    };
+    let len = rest.partition_point(|item| {
+        let item_ptr = core::ptr::addr_of!(*item);
+        // SAFETY: `item` is guaranteed to be in bounds of `slice`.
+        let elem_distance = usize::try_from(unsafe { item_ptr.offset_from(first) }).unwrap();
+        let value_distance = item.checked_sub(*first).expect("array must be sorted");
+        elem_distance == usize::from(value_distance)
+    });
+    len + 1 // +1 for the first element
+}
+
+/// Return the last contiguous range of elements in a sorted slice.
+pub(crate) fn last_contiguous_range_len(slice: &[u16]) -> usize {
+    let [rest @ .., last] = slice else {
+        // Explicitly empty range
+        return 0;
+    };
+    let last_ptr = core::ptr::addr_of!(*last);
+    let len_from_start = rest.partition_point(|item| {
+        // SAFETY: `item` is guaranteed to be in bounds of `slice`.
+        let elem_distance = usize::try_from(unsafe { last_ptr.offset_from(item) }).unwrap();
+        let value_distance = last.checked_sub(*item).expect("array must be sorted");
+        elem_distance != usize::from(value_distance)
+    });
+    slice.len() - len_from_start
+}
+
 impl ArrayStore {
     pub fn new() -> ArrayStore {
         ArrayStore { vec: vec![] }