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()`.

Co-authored-by: Matthew Herzl <[email protected]>
  • Loading branch information
@Dr-Emann @mherzl
Dr-Emann and mherzl committed Nov 14, 2024
1 parent dd8bc5a commit d6ad8c4
Show file tree
Hide file tree
Showing 4 changed files with 279 additions and 38 deletions.
20 changes: 10 additions & 10 deletions roaring/src/bitmap/inherent.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,8 +93,8 @@ impl RoaringBitmap {
R: RangeBounds<u32>,
{
let (start, end) = match util::convert_range_to_inclusive(range) {
Some(range) => (*range.start(), *range.end()),
None => return 0,
Ok(range) => (*range.start(), *range.end()),
Err(_) => return 0,
};

let (start_container_key, start_index) = util::split(start);
Expand Down Expand Up @@ -238,8 +238,8 @@ impl RoaringBitmap {
R: RangeBounds<u32>,
{
let (start, end) = match util::convert_range_to_inclusive(range) {
Some(range) => (*range.start(), *range.end()),
None => return 0,
Ok(range) => (*range.start(), *range.end()),
Err(_) => return 0,
};

let (start_container_key, start_index) = util::split(start);
Expand Down Expand Up @@ -308,9 +308,9 @@ impl RoaringBitmap {
R: RangeBounds<u32>,
{
let (start, end) = match util::convert_range_to_inclusive(range) {
Some(range) => (*range.start(), *range.end()),
// Empty ranges are always contained
None => return true,
Ok(range) => (*range.start(), *range.end()),
// Empty/Invalid ranges are always contained
Err(_) => return true,
};
let (start_high, start_low) = util::split(start);
let (end_high, end_low) = util::split(end);
Expand Down Expand Up @@ -368,9 +368,9 @@ impl RoaringBitmap {
R: RangeBounds<u32>,
{
let (start, end) = match util::convert_range_to_inclusive(range) {
Some(range) => (*range.start(), *range.end()),
// Empty ranges have 0 bits set in them
None => return 0,
Ok(range) => (*range.start(), *range.end()),
// Empty/invalid ranges have 0 bits set in them
Err(_) => return 0,
};

let (start_key, start_low) = util::split(start);
Expand Down
120 changes: 120 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,117 @@ 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.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.range(100..);
/// assert_eq!(iter.next(), Some(u32::MAX));
/// assert_eq!(iter.next(), None);
///
/// let mut iter = bitmap.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 range<R>(&self, range: R) -> Iter<'_>
where
R: RangeBounds<u32>,
{
let range = match util::convert_range_to_inclusive(range) {
Ok(range) => range,
Err(util::ConvertRangeError::Empty) => return Iter::empty(),
Err(util::ConvertRangeError::StartGreaterThanEnd) => {
panic!("range start is greater than range end")
}
Err(util::ConvertRangeError::StartAndEndEqualExcluded) => {
panic!("range start and end are equal and excluded")
}
};
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_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_range(100..);
/// assert_eq!(iter.next(), Some(u32::MAX));
/// assert_eq!(iter.next(), None);
///
/// let mut iter = bitmap().into_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_range<R>(self, range: R) -> IntoIter
where
R: RangeBounds<u32>,
{
let range = match util::convert_range_to_inclusive(range) {
Ok(range) => range,
Err(util::ConvertRangeError::Empty) => return IntoIter::empty(),
Err(util::ConvertRangeError::StartGreaterThanEnd) => {
panic!("range start is greater than range end")
}
Err(util::ConvertRangeError::StartAndEndEqualExcluded) => {
panic!("range start and end are equal and excluded")
}
};
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
82 changes: 54 additions & 28 deletions roaring/src/bitmap/util.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,30 +14,56 @@ pub fn join(high: u16, low: u16) -> u32 {
(u32::from(high) << 16) + u32::from(low)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ConvertRangeError {
Empty,
StartGreaterThanEnd,
StartAndEndEqualExcluded,
}

/// Convert a `RangeBounds<u32>` object to `RangeInclusive<u32>`,
pub fn convert_range_to_inclusive<R>(range: R) -> Option<RangeInclusive<u32>>
pub fn convert_range_to_inclusive<R>(range: R) -> Result<RangeInclusive<u32>, ConvertRangeError>
where
R: RangeBounds<u32>,
{
let start: u32 = match range.start_bound() {
Bound::Included(&i) => i,
Bound::Excluded(&i) => i.checked_add(1)?,
Bound::Unbounded => 0,
};
let end: u32 = match range.end_bound() {
Bound::Included(&i) => i,
Bound::Excluded(&i) => i.checked_sub(1)?,
Bound::Unbounded => u32::MAX,
};
if end < start {
return None;
let start_bound = range.start_bound().cloned();
let end_bound = range.end_bound().cloned();
match (start_bound, end_bound) {
(Bound::Excluded(s), Bound::Excluded(e)) if s == e => {
Err(ConvertRangeError::StartAndEndEqualExcluded)
}
(Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e))
if s > e =>
{
Err(ConvertRangeError::StartGreaterThanEnd)
}
_ => {
let start = match start_bound {
Bound::Included(s) => s,
Bound::Excluded(s) => s.checked_add(1).ok_or(ConvertRangeError::Empty)?,
Bound::Unbounded => 0,
};

let end = match end_bound {
Bound::Included(e) => e,
Bound::Excluded(e) => e.checked_sub(1).ok_or(ConvertRangeError::Empty)?,
Bound::Unbounded => u32::MAX,
};

if start > end {
// This handles e.g. `x..x`: we've ruled out `start > end` overall, so a value must
// have been changed via exclusion.
Err(ConvertRangeError::Empty)
} else {
Ok(start..=end)
}
}
}
Some(start..=end)
}

#[cfg(test)]
mod test {
use super::{convert_range_to_inclusive, join, split};
use super::{convert_range_to_inclusive, join, split, ConvertRangeError};
use core::ops::Bound;

#[test]
Expand Down Expand Up @@ -67,27 +93,27 @@ mod test {
#[test]
#[allow(clippy::reversed_empty_ranges)]
fn test_convert_range_to_inclusive() {
assert_eq!(Some(1..=5), convert_range_to_inclusive(1..6));
assert_eq!(Some(1..=u32::MAX), convert_range_to_inclusive(1..));
assert_eq!(Some(0..=u32::MAX), convert_range_to_inclusive(..));
assert_eq!(Some(16..=16), convert_range_to_inclusive(16..=16));
assert_eq!(Ok(1..=5), convert_range_to_inclusive(1..6));
assert_eq!(Ok(1..=u32::MAX), convert_range_to_inclusive(1..));
assert_eq!(Ok(0..=u32::MAX), convert_range_to_inclusive(..));
assert_eq!(Ok(16..=16), convert_range_to_inclusive(16..=16));
assert_eq!(
Some(11..=19),
Ok(11..=19),
convert_range_to_inclusive((Bound::Excluded(10), Bound::Excluded(20)))
);

assert_eq!(None, convert_range_to_inclusive(0..0));
assert_eq!(None, convert_range_to_inclusive(5..5));
assert_eq!(None, convert_range_to_inclusive(1..0));
assert_eq!(None, convert_range_to_inclusive(10..5));
assert_eq!(Err(ConvertRangeError::Empty), convert_range_to_inclusive(0..0));
assert_eq!(Err(ConvertRangeError::Empty), convert_range_to_inclusive(5..5));
assert_eq!(Err(ConvertRangeError::StartGreaterThanEnd), convert_range_to_inclusive(1..0));
assert_eq!(Err(ConvertRangeError::StartGreaterThanEnd), convert_range_to_inclusive(10..5));
assert_eq!(
None,
Err(ConvertRangeError::Empty),
convert_range_to_inclusive((Bound::Excluded(u32::MAX), Bound::Included(u32::MAX)))
);
assert_eq!(
None,
convert_range_to_inclusive((Bound::Excluded(u32::MAX), Bound::Included(u32::MAX)))
Err(ConvertRangeError::StartAndEndEqualExcluded),
convert_range_to_inclusive((Bound::Excluded(u32::MAX), Bound::Excluded(u32::MAX)))
);
assert_eq!(None, convert_range_to_inclusive((Bound::Excluded(0), Bound::Included(0))));
assert_eq!(Err(ConvertRangeError::Empty), convert_range_to_inclusive((Bound::Excluded(0), Bound::Included(0))));
}
}
Loading

0 comments on commit d6ad8c4

Please sign in to comment.