diff --git a/roaring/src/bitmap/iter.rs b/roaring/src/bitmap/iter.rs index de14092a..766d7ea6 100644 --- a/roaring/src/bitmap/iter.rs +++ b/roaring/src/bitmap/iter.rs @@ -1,5 +1,6 @@ use alloc::vec; use core::iter::FusedIterator; +use core::ops::RangeBounds; use core::slice; use super::container::Container; @@ -153,6 +154,10 @@ impl Iter<'_> { Iter { front: None, containers: containers.iter(), back: None } } + fn empty() -> Self { + Self::new(&[]) + } + /// Advance the iterator to the first position where the item has a value >= `n` /// /// # Examples @@ -197,6 +202,10 @@ impl IntoIter { IntoIter { front: None, containers: containers.into_iter(), back: None } } + fn empty() -> Self { + Self::new(Vec::new()) + } + /// Advance the iterator to the first position where the item has a value >= `n` /// /// # Examples @@ -550,6 +559,105 @@ impl RoaringBitmap { pub fn iter(&self) -> Iter { Iter::new(&self.containers) } + + /// Iterator over values within a range stored in the RoaringBitmap. + /// + /// # Examples + /// + /// ```rust + /// use core::ops::Bound; + /// use roaring::RoaringBitmap; + /// + /// let bitmap = RoaringBitmap::from([0, 1, 2, 3, 4, 5, 10, 11, 12, 20, 21, u32::MAX]); + /// let mut iter = bitmap.iter_range(10..20); + /// + /// assert_eq!(iter.next(), Some(10)); + /// assert_eq!(iter.next(), Some(11)); + /// assert_eq!(iter.next(), Some(12)); + /// assert_eq!(iter.next(), None); + /// + /// let mut iter = bitmap.iter_range(100..); + /// assert_eq!(iter.next(), Some(u32::MAX)); + /// assert_eq!(iter.next(), None); + /// + /// let mut iter = bitmap.iter_range((Bound::Excluded(0), Bound::Included(10))); + /// assert_eq!(iter.next(), Some(1)); + /// assert_eq!(iter.next(), Some(2)); + /// assert_eq!(iter.next(), Some(3)); + /// assert_eq!(iter.next(), Some(4)); + /// assert_eq!(iter.next(), Some(5)); + /// assert_eq!(iter.next(), Some(10)); + /// assert_eq!(iter.next(), None); + /// ``` + pub fn iter_range(&self, range: R) -> Iter<'_> + where + R: RangeBounds, + { + let range = match util::convert_range_to_inclusive(range) { + Some(range) => range, + None => return Iter::empty(), + }; + let (start, end) = (*range.start(), *range.end()); + let mut iter = self.iter(); + if start != 0 { + iter.advance_to(start); + } + if end != u32::MAX { + iter.advance_back_to(end); + } + iter + } + + /// Iterator over values within a range stored in the RoaringBitmap. + /// + /// # Examples + /// + /// ```rust + /// use core::ops::Bound; + /// use roaring::RoaringBitmap; + /// + /// fn bitmap() -> RoaringBitmap { + /// RoaringBitmap::from([0, 1, 2, 3, 4, 5, 10, 11, 12, 20, 21, u32::MAX]) + /// } + /// + /// let mut iter = bitmap().into_iter_range(10..20); + /// + /// assert_eq!(iter.next(), Some(10)); + /// assert_eq!(iter.next(), Some(11)); + /// assert_eq!(iter.next(), Some(12)); + /// assert_eq!(iter.next(), None); + /// + /// let mut iter = bitmap().into_iter_range(100..); + /// assert_eq!(iter.next(), Some(u32::MAX)); + /// assert_eq!(iter.next(), None); + /// + /// let mut iter = bitmap().into_iter_range((Bound::Excluded(0), Bound::Included(10))); + /// assert_eq!(iter.next(), Some(1)); + /// assert_eq!(iter.next(), Some(2)); + /// assert_eq!(iter.next(), Some(3)); + /// assert_eq!(iter.next(), Some(4)); + /// assert_eq!(iter.next(), Some(5)); + /// assert_eq!(iter.next(), Some(10)); + /// assert_eq!(iter.next(), None); + /// ``` + pub fn into_iter_range(self, range: R) -> IntoIter + where + R: RangeBounds, + { + let range = match util::convert_range_to_inclusive(range) { + Some(range) => range, + None => return IntoIter::empty(), + }; + let (start, end) = (*range.start(), *range.end()); + let mut iter = self.into_iter(); + if start != 0 { + iter.advance_to(start); + } + if end != u32::MAX { + iter.advance_back_to(end); + } + iter + } } impl<'a> IntoIterator for &'a RoaringBitmap { diff --git a/roaring/tests/iter_range.rs b/roaring/tests/iter_range.rs new file mode 100644 index 00000000..75d54c7b --- /dev/null +++ b/roaring/tests/iter_range.rs @@ -0,0 +1,54 @@ +use proptest::collection::btree_set; +use proptest::prelude::*; +use roaring::RoaringBitmap; + +#[test] +fn range_array() { + let mut rb = RoaringBitmap::new(); + rb.insert(0); + rb.insert(1); + rb.insert(10); + rb.insert(100_000); + rb.insert(999_999); + rb.insert(1_000_000); + + let expected = vec![1, 10, 100_000, 999_999]; + let actual: Vec = rb.iter_range(1..=999_999).collect(); + assert_eq!(expected, actual); +} + +#[test] +fn range_bitmap() { + let rb = RoaringBitmap::from_sorted_iter(10..5000).unwrap(); + + let expected = vec![10, 11, 12]; + let actual: Vec = rb.iter_range(0..13).collect(); + assert_eq!(expected, actual); +} + +#[test] +fn range_none() { + let rb = RoaringBitmap::from_sorted_iter(10..5000).unwrap(); + + let expected: Vec = vec![]; + #[allow(clippy::reversed_empty_ranges)] + let actual: Vec = rb.iter_range(13..0).collect(); + assert_eq!(expected, actual); +} + +proptest! { + #[test] + fn proptest_range( + values in btree_set(..=262_143_u32, ..=1000), + range_a in 0u32..262_143, + range_b in 0u32..262_143, + ){ + let range = range_a.min(range_b)..=range_a.max(range_b); + + let bitmap = RoaringBitmap::from_sorted_iter(values.iter().cloned()).unwrap(); + let expected: Vec = values.range(range.clone()).copied().collect(); + let actual: Vec = bitmap.iter_range(range.clone()).collect(); + + assert_eq!(expected, actual); + } +}