Skip to content

Commit

Permalink
implement range based iteration
Browse files Browse the repository at this point in the history 
Implement `bitmap.iter_range()` and `bitmap.into_iter_range()`, based on
`advance_to()` and `advance_back_to()`.

Open questions:
- The equivalent to `.iter_range` for BTreeSet is called `.range`. Is that a
  better name: `.range`/`.into_range`?
- `.range` panics on start > end. (or start == end when both are excluded).
  Should we?

Co-authored-by: Matthew Herzl <[email protected]>
  • Loading branch information
@Dr-Emann @mherzl
Dr-Emann and mherzl committed Nov 11, 2024
1 parent 9587082 commit 03682b5
Show file tree
Hide file tree
Showing 2 changed files with 162 additions and 0 deletions.
108 changes: 108 additions & 0 deletions roaring/src/bitmap/iter.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use alloc::vec;
use core::iter::FusedIterator;
use core::ops::RangeBounds;
use core::slice;

use super::container::Container;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<R>(&self, range: R) -> Iter<'_>
where
R: RangeBounds<u32>,
{
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<R>(self, range: R) -> IntoIter
where
R: RangeBounds<u32>,
{
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 {
Expand Down
54 changes: 54 additions & 0 deletions roaring/tests/iter_range.rs
View file
Original file line number Diff line number Diff line change
@@ -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<u32> = 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<u32> = 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<u32> = vec![];
#[allow(clippy::reversed_empty_ranges)]
let actual: Vec<u32> = 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<u32> = values.range(range.clone()).copied().collect();
let actual: Vec<u32> = bitmap.iter_range(range.clone()).collect();

assert_eq!(expected, actual);
}
}

0 comments on commit 03682b5

Please sign in to comment.