diff --git a/roaring/src/bitmap/container.rs b/roaring/src/bitmap/container.rs index 3ab1e5ad..0af1a558 100644 --- a/roaring/src/bitmap/container.rs +++ b/roaring/src/bitmap/container.rs @@ -300,6 +300,10 @@ impl<'a> Iterator for Iter<'a> { fn next(&mut self) -> Option { self.inner.next().map(|i| util::join(self.key, i)) } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } } impl DoubleEndedIterator for Iter<'_> { @@ -313,3 +317,15 @@ impl fmt::Debug for Container { format!("Container<{:?} @ {:?}>", self.len(), self.key).fmt(formatter) } } + +impl<'a> Iter<'a> { + pub fn empty() -> Self { + Self { key: 0, inner: store::Iter::Empty } + } +} + +impl AsRef for Container { + fn as_ref(&self) -> &Container { + self + } +} diff --git a/roaring/src/bitmap/iter.rs b/roaring/src/bitmap/iter.rs index 59463a21..a63c0241 100644 --- a/roaring/src/bitmap/iter.rs +++ b/roaring/src/bitmap/iter.rs @@ -1,132 +1,199 @@ use alloc::vec; -use core::iter; use core::slice; use super::container::Container; use crate::{NonSortedIntegers, RoaringBitmap}; +use crate::bitmap::{container, util}; #[cfg(not(feature = "std"))] use alloc::vec::Vec; +use std::iter::Peekable; + +macro_rules! impl_iter { + ($t:ident,$c:ty, $c_iter:ty, $iter_lt:lifetime $(,$lt:lifetime)?) => { + /// An iterator for `RoaringBitmap`. + pub struct $t$(<$lt>)? { + containers_iter: $c_iter, + iter_front: Option>, + iter_back: Option>, + size_hint: u64, + } -/// An iterator for `RoaringBitmap`. -pub struct Iter<'a> { - inner: iter::Flatten>, - size_hint: u64, -} - -/// An iterator for `RoaringBitmap`. -pub struct IntoIter { - inner: iter::Flatten>, - size_hint: u64, -} - -impl Iter<'_> { - fn new(containers: &[Container]) -> Iter { - let size_hint = containers.iter().map(|c| c.len()).sum(); - Iter { inner: containers.iter().flatten(), size_hint } - } -} - -impl IntoIter { - fn new(containers: Vec) -> IntoIter { - let size_hint = containers.iter().map(|c| c.len()).sum(); - IntoIter { inner: containers.into_iter().flatten(), size_hint } - } -} + impl$(<$lt>)? $t$(<$lt>)? { + fn new(containers: $c) -> Self { + let size_hint = containers.iter().map(|c| c.len()).sum(); + let containers_iter = containers.into_iter(); + Self { + containers_iter, + iter_front: None, + iter_back: None, + size_hint, + } + } + } -impl Iterator for Iter<'_> { - type Item = u32; + impl$(<$lt>)? Iterator for $t$(<$lt>)? { + type Item = u32; + + fn next(&mut self) -> Option { + self.size_hint = self.size_hint.saturating_sub(1); + loop { + if let Some(iter) = &mut self.iter_front { + if let item @ Some(_) = iter.next() { + return item + } + } + if let Some(container) = self.containers_iter.next() { + self.iter_front = Some(container.into_iter()) + } else if let Some(iter) = &mut self.iter_back { + return iter.next() + } else { + return None + } + } + } - fn next(&mut self) -> Option { - self.size_hint = self.size_hint.saturating_sub(1); - self.inner.next() - } + fn size_hint(&self) -> (usize, Option) { + if self.size_hint < usize::MAX as u64 { + (self.size_hint as usize, Some(self.size_hint as usize)) + } else { + (usize::MAX, None) + } + } - fn size_hint(&self) -> (usize, Option) { - if self.size_hint < usize::MAX as u64 { - (self.size_hint as usize, Some(self.size_hint as usize)) - } else { - (usize::MAX, None) + #[inline] + fn fold(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + match (self.iter_front, self.iter_back) { + (Some(iter_front), Some(iter_back)) => { + iter_front.chain(self.containers_iter.flatten()).chain(iter_back).fold(init, f) + }, + (Some(iter_front), None) => { + iter_front.chain(self.containers_iter.flatten()).fold(init, f) + }, + (None, Some(iter_back)) => { + self.containers_iter.flatten().chain(iter_back).fold(init, f) + }, + (None, None) => self.containers_iter.flatten().fold(init, f) + } + } } - } - #[inline] - fn fold(self, init: B, f: F) -> B - where - Self: Sized, - F: FnMut(B, Self::Item) -> B, - { - self.inner.fold(init, f) - } -} - -impl DoubleEndedIterator for Iter<'_> { - fn next_back(&mut self) -> Option { - self.size_hint = self.size_hint.saturating_sub(1); - self.inner.next_back() - } + impl$(<$lt>)? DoubleEndedIterator for $t$(<$lt>)? { + fn next_back(&mut self) -> Option { + self.size_hint = self.size_hint.saturating_sub(1); + loop { + if let Some(iter) = &mut self.iter_back { + if let item @ Some(_) = iter.next_back() { + return item + } + } + if let Some(container) = self.containers_iter.next_back() { + self.iter_back = Some(container.into_iter()) + } else if let Some(iter) = &mut self.iter_front { + return iter.next_back() + } else { + return None + } + } + } - #[inline] - fn rfold(self, init: Acc, fold: Fold) -> Acc - where - Fold: FnMut(Acc, Self::Item) -> Acc, - { - self.inner.rfold(init, fold) - } + #[inline] + fn rfold(self, init: Acc, f: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + match (self.iter_front, self.iter_back) { + (Some(iter_front), Some(iter_back)) => { + iter_front.chain(self.containers_iter.flatten()).chain(iter_back).rfold(init, f) + }, + (Some(iter_front), None) => { + iter_front.chain(self.containers_iter.flatten()).rfold(init, f) + }, + (None, Some(iter_back)) => { + self.containers_iter.flatten().chain(iter_back).rfold(init, f) + }, + (None, None) => self.containers_iter.flatten().rfold(init, f) + } + } + } + #[cfg(target_pointer_width = "64")] + impl$(<$lt>)? ExactSizeIterator for $t$(<$lt>)? { + fn len(&self) -> usize { + self.size_hint as usize + } + } + }; } +impl_iter!(Iter, &'a [Container], slice::Iter<'a, Container>, 'a, 'a); +impl_iter!(IntoIter, Vec, vec::IntoIter, 'static); -#[cfg(target_pointer_width = "64")] -impl ExactSizeIterator for Iter<'_> { - fn len(&self) -> usize { - self.size_hint as usize - } +pub struct AdvanceToIter<'a, CI> { + containers_iter: CI, + iter_front: Peekable>, + iter_back: Option>>, } -impl Iterator for IntoIter { - type Item = u32; - - fn next(&mut self) -> Option { - self.size_hint = self.size_hint.saturating_sub(1); - self.inner.next() - } - - fn size_hint(&self) -> (usize, Option) { - if self.size_hint < usize::MAX as u64 { - (self.size_hint as usize, Some(self.size_hint as usize)) +impl<'a, CI> AdvanceToIter<'a, CI> +where + Self: AdvanceIterContainer<'a>, +{ + fn new( + containers_iter: CI, + iter_front: Option>, + iter_back: Option>, + n: u32, + ) -> Self { + let mut result = Self { + containers_iter, + iter_front: container::Iter::empty().peekable(), + iter_back: iter_back.map(|o| o.peekable()), + }; + if let Some(iter_front) = iter_front { + result.iter_front = iter_front.peekable(); } else { - (usize::MAX, None) + result.advance_container(); } - } - - #[inline] - fn fold(self, init: B, f: F) -> B - where - Self: Sized, - F: FnMut(B, Self::Item) -> B, - { - self.inner.fold(init, f) + if let Some(peek) = result.iter_front.peek().cloned() { + if peek < n { + let (peek_key, _) = util::split(peek); + let (target_key, _) = util::split(n); + if target_key > peek_key { + while let Some(next_key) = result.advance_container() { + if next_key >= target_key { + break; + } + } + } + while let Some(peek) = result.iter_front.peek() { + if *peek >= n { + break; + } else { + result.iter_front.next(); + } + } + } + } + result } } -impl DoubleEndedIterator for IntoIter { - fn next_back(&mut self) -> Option { - self.size_hint = self.size_hint.saturating_sub(1); - self.inner.next_back() - } - - #[inline] - fn rfold(self, init: Acc, fold: Fold) -> Acc - where - Fold: FnMut(Acc, Self::Item) -> Acc, - { - self.inner.rfold(init, fold) - } -} +impl<'a, CI> Iterator for AdvanceToIter<'a, CI> +where + Self: AdvanceIterContainer<'a>, +{ + type Item = u32; -#[cfg(target_pointer_width = "64")] -impl ExactSizeIterator for IntoIter { - fn len(&self) -> usize { - self.size_hint as usize + fn next(&mut self) -> Option { + loop { + if let item @ Some(_) = self.iter_front.next() { + return item; + } + self.advance_container()?; + } } } @@ -149,6 +216,41 @@ impl RoaringBitmap { pub fn iter(&self) -> Iter { Iter::new(&self.containers) } + + /// Iterator over each value >= `n` stored in the RoaringBitmap, guarantees values are ordered + /// by value. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use core::iter::FromIterator; + /// + /// let bitmap = (1..3).collect::(); + /// let mut iter = bitmap.iter_from(2); + /// + /// assert_eq!(iter.next(), Some(2)); + /// assert_eq!(iter.next(), None); + /// ``` + pub fn iter_from(&self, n: u32) -> AdvanceToIter<'_, slice::Iter<'_, Container>> { + let (key, _) = util::split(n); + match self.containers.binary_search_by_key(&key, |container| container.key) { + Ok(index) | Err(index) => { + if index == self.containers.len() { + // no container has a key >= key(n) + AdvanceToIter::new([].iter(), None, None, n) + } else { + let containers = &self.containers[index..]; + let iter = (&containers[0]).into_iter(); + if index + 1 < containers.len() - 1 { + AdvanceToIter::new(containers[index + 1..].iter(), Some(iter), None, n) + } else { + AdvanceToIter::new([].iter(), Some(iter), None, n) + } + } + } + } + } } impl<'a> IntoIterator for &'a RoaringBitmap { @@ -298,3 +400,55 @@ impl RoaringBitmap { Ok(count) } } + +pub trait AdvanceIterContainer<'a> { + fn advance_container(&mut self) -> Option; +} + +macro_rules! impl_advance_iter_container { + ($lt:lifetime,$ty:ty) => { + impl<$lt> AdvanceIterContainer<$lt> for AdvanceToIter<$lt, $ty> { + fn advance_container(&mut self) -> Option { + if let Some(container) = self.containers_iter.next() { + let result = container.key; + self.iter_front = container.into_iter().peekable(); + Some(result) + } else if let Some(iter_back) = &mut self.iter_back { + std::mem::swap(iter_back, &mut self.iter_front); + self.iter_back = None; + if let Some(v) = self.iter_front.peek().cloned() { + let (key, _) = util::split(v); + Some(key) + } else { + None + } + } else { + None + } + } + } + }; +} +impl_advance_iter_container!('a, slice::Iter<'a, Container>); +impl_advance_iter_container!('a ,vec::IntoIter); + +macro_rules! impl_advance_to { + ($ty:ty, $ret:ty $(,$lt:lifetime)? ) => { + impl$(<$lt>)? $ty { + /// Advance the iterator to the first position where the item has a value >= `n` + pub fn advance_to(mut self, n: u32) -> $ret { + if let Some(iter_front) = self.iter_front { + AdvanceToIter::new(self.containers_iter, Some(iter_front), self.iter_back, n) + } else { + if let Some(container) = self.containers_iter.next() { + AdvanceToIter::new(self.containers_iter, Some(container.into_iter()), self.iter_back, n) + } else { + AdvanceToIter::new(self.containers_iter, Some(container::Iter::empty()), self.iter_back, n) + } + } + } + } + }; +} +impl_advance_to!(Iter<'a>, AdvanceToIter<'a, slice::Iter<'a, Container>>, 'a); +impl_advance_to!(IntoIter, AdvanceToIter<'static, vec::IntoIter>); diff --git a/roaring/src/bitmap/store/array_store/mod.rs b/roaring/src/bitmap/store/array_store/mod.rs index 883db31f..541b38f6 100644 --- a/roaring/src/bitmap/store/array_store/mod.rs +++ b/roaring/src/bitmap/store/array_store/mod.rs @@ -11,11 +11,10 @@ use core::ops::{BitAnd, BitAndAssign, BitOr, BitXor, RangeInclusive, Sub, SubAss #[cfg(not(feature = "std"))] use alloc::vec::Vec; +use super::bitmap_store::{bit, key, BitmapStore, BITMAP_LENGTH}; #[cfg(not(feature = "std"))] use alloc::boxed::Box; -use super::bitmap_store::{bit, key, BitmapStore, BITMAP_LENGTH}; - #[derive(Clone, Eq, PartialEq)] pub struct ArrayStore { vec: Vec, diff --git a/roaring/src/bitmap/store/mod.rs b/roaring/src/bitmap/store/mod.rs index bb0d5822..f1372f19 100644 --- a/roaring/src/bitmap/store/mod.rs +++ b/roaring/src/bitmap/store/mod.rs @@ -30,6 +30,7 @@ pub enum Iter<'a> { Vec(vec::IntoIter), BitmapBorrowed(BitmapIter<&'a [u64; BITMAP_LENGTH]>), BitmapOwned(BitmapIter>), + Empty, } impl Store { @@ -506,6 +507,7 @@ impl<'a> Iterator for Iter<'a> { Iter::Vec(inner) => inner.next(), Iter::BitmapBorrowed(inner) => inner.next(), Iter::BitmapOwned(inner) => inner.next(), + Iter::Empty => None, } } } @@ -517,6 +519,7 @@ impl DoubleEndedIterator for Iter<'_> { Iter::Vec(inner) => inner.next_back(), Iter::BitmapBorrowed(inner) => inner.next_back(), Iter::BitmapOwned(inner) => inner.next_back(), + Iter::Empty => None, } } } diff --git a/roaring/src/lib.rs b/roaring/src/lib.rs index b44c19c2..b4cbf5a5 100644 --- a/roaring/src/lib.rs +++ b/roaring/src/lib.rs @@ -19,6 +19,7 @@ extern crate byteorder; #[macro_use] extern crate alloc; +extern crate core; use core::fmt; diff --git a/roaring/tests/iter_advance_to.rs b/roaring/tests/iter_advance_to.rs new file mode 100644 index 00000000..7aecc3d7 --- /dev/null +++ b/roaring/tests/iter_advance_to.rs @@ -0,0 +1,79 @@ +use roaring::RoaringBitmap; + +#[test] +fn iter_basic() { + let bm = RoaringBitmap::from([1, 2, 3, 4, 11, 12, 13, 14]); + let mut i = bm.iter().advance_to(10); + for n in 11..=14 { + assert_eq!(i.next(), Some(n)) + } + assert_eq!(i.next(), None); +} + +#[test] +fn iter_advance_past_end() { + let bm = RoaringBitmap::from([1, 2, 3, 4, 11, 12, 13, 14]); + let mut i = bm.iter().advance_to(15); + assert_eq!(i.next(), None); +} + +#[test] +fn iter_multi_container() { + let bm = RoaringBitmap::from([1, 2, 3, 100000, 100001]); + let mut i = bm.iter().advance_to(3); + assert_eq!(i.next(), Some(3)); + assert_eq!(i.next(), Some(100000)); + assert_eq!(i.next(), Some(100001)); + assert_eq!(i.next(), None); +} + +#[test] +fn iter_empty() { + let bm = RoaringBitmap::new(); + assert_eq!(bm.iter().advance_to(31337).next(), None) +} + +#[test] +fn into_iter_basic() { + let bm = RoaringBitmap::from([1, 2, 3, 4, 11, 12, 13, 14]); + let mut i = bm.into_iter().advance_to(10); + for n in 11..=14 { + assert_eq!(i.next(), Some(n)) + } + assert_eq!(i.next(), None); +} + +#[test] +fn into_iter_multi_container() { + let bm = RoaringBitmap::from([1, 2, 3, 100000, 100001]); + let mut i = bm.into_iter().advance_to(3); + assert_eq!(i.next(), Some(3)); + assert_eq!(i.next(), Some(100000)); + assert_eq!(i.next(), Some(100001)); + assert_eq!(i.next(), None); +} + +#[test] +fn into_iter_empty() { + let bm = RoaringBitmap::new(); + assert_eq!(bm.into_iter().advance_to(31337).next(), None) +} + +#[test] +fn iter_from() { + let bm = RoaringBitmap::from([1, 2, 3, 100000, 100001]); + let mut i = bm.iter_from(99999); + assert_eq!(i.next(), Some(100000)); + assert_eq!(i.next(), Some(100001)); + assert_eq!(i.next(), None); +} + +#[test] +fn advance_to_with_tail_iter() { + let bm = RoaringBitmap::from([1, 2, 3, 100000, 100001]); + let mut i = bm.iter(); + i.next_back(); + let mut i = i.advance_to(100000); + assert_eq!(i.next(), Some(100000)); + assert_eq!(i.next(), None); +}