diff --git a/roaring/src/bitmap/store/array_store/mod.rs b/roaring/src/bitmap/store/array_store/mod.rs index 61967704..ec3cbcf9 100644 --- a/roaring/src/bitmap/store/array_store/mod.rs +++ b/roaring/src/bitmap/store/array_store/mod.rs @@ -102,11 +102,10 @@ impl ArrayStore { let mut vec = Vec::with_capacity(bits_set as usize); - let chunks = bytes.chunks_exact(size_of::()); - let remainder = chunks.remainder(); - for (index, chunk) in chunks.enumerate() { + let (chunks, remainder) = bytes.as_chunks::<{ size_of::() }>(); + for (index, chunk) in chunks.iter().enumerate() { let bit_index = (byte_offset + index * size_of::()) * 8; - let mut word = Word::from_le_bytes(chunk.try_into().unwrap()); + let mut word = Word::from_le_bytes(*chunk); while word != 0 { vec.push((word.trailing_zeros() + bit_index as u32) as u16); diff --git a/roaring/src/bitmap/store/mod.rs b/roaring/src/bitmap/store/mod.rs index 11a3de0c..c96e6045 100644 --- a/roaring/src/bitmap/store/mod.rs +++ b/roaring/src/bitmap/store/mod.rs @@ -66,10 +66,9 @@ impl Store { // using u64s than for each byte let bits_set = { let mut bits_set = 0; - let chunks = bytes.chunks_exact(mem::size_of::()); - let remainder = chunks.remainder(); + let (chunks, remainder) = bytes.as_chunks::<{ mem::size_of::() }>(); for chunk in chunks { - let chunk = u64::from_ne_bytes(chunk.try_into().unwrap()); + let chunk = u64::from_ne_bytes(*chunk); bits_set += u64::from(chunk.count_ones()); } for byte in remainder { diff --git a/roaring/src/lib.rs b/roaring/src/lib.rs index f6c4453b..d66f15f5 100644 --- a/roaring/src/lib.rs +++ b/roaring/src/lib.rs @@ -43,6 +43,9 @@ impl fmt::Display for IntegerTooSmall { } } +#[cfg(feature = "std")] +impl std::error::Error for IntegerTooSmall {} + /// An error type that is returned when an iterator isn't sorted. #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct NonSortedIntegers { diff --git a/roaring/src/treemap/iter.rs b/roaring/src/treemap/iter.rs index 36a0bc16..f7f7bae1 100644 --- a/roaring/src/treemap/iter.rs +++ b/roaring/src/treemap/iter.rs @@ -1,4 +1,5 @@ use alloc::collections::{btree_map, BTreeMap}; +use core::cmp::Ordering; use core::iter; use core::ops::Add; @@ -147,24 +148,39 @@ impl Iter<'_> { pub fn advance_to(&mut self, n: u64) { let (key, index) = util::split(n); + if let Some(ref mut front) = self.front { + match front.hi.cmp(&key) { + Ordering::Less => {} + Ordering::Equal => { + front.advance_to(index); + return; + } + Ordering::Greater => return, + } + self.front = None; + } + self.outer.advance_to(key); - if self.front.is_none() { - let Some(next) = self.outer.next() else { - // if the current front iterator is empty or not yet initialized, - // but the outer bitmap iterator is empty, then consume the back - // iterator from the front if it is not also exhausted - if let Some(ref mut back) = self.back { - back.advance_to(index); + let Some(next) = self.outer.next() else { + // if the current front iterator is empty or not yet initialized, + // but the outer bitmap iterator is empty, then consume the back + // iterator from the front if it is not also exhausted + if let Some(ref mut back) = self.back { + match back.hi.cmp(&key) { + Ordering::Less => self.back = None, + Ordering::Equal => back.advance_to(index), + Ordering::Greater => {} } - return; - }; - self.front = Some(to64iter(next)); - } + } + return; + }; - if let Some(ref mut front) = self.front { + let mut front = to64iter(next); + if front.hi == key { front.advance_to(index); } + self.front = Some(front); } /// Advance the back of the iterator to the first position where the item has a value <= `n` @@ -185,24 +201,39 @@ impl Iter<'_> { pub fn advance_back_to(&mut self, n: u64) { let (key, index) = util::split(n); + if let Some(ref mut back) = self.back { + match back.hi.cmp(&key) { + Ordering::Less => return, + Ordering::Equal => { + back.advance_back_to(index); + return; + } + Ordering::Greater => {} + } + self.back = None; + } + self.outer.advance_back_to(key); - if self.back.is_none() { - let Some(next_back) = self.outer.next_back() else { - // if the current back iterator is empty or not yet initialized, - // but the outer bitmap iterator is empty, then consume the front - // iterator from the back if it is not also exhausted - if let Some(ref mut front) = self.front { - front.advance_back_to(index); + let Some(next_back) = self.outer.next_back() else { + // if the current back iterator is empty or not yet initialized, + // but the outer bitmap iterator is empty, then consume the front + // iterator from the back if it is not also exhausted + if let Some(ref mut front) = self.front { + match front.hi.cmp(&key) { + Ordering::Less => {} + Ordering::Equal => front.advance_back_to(index), + Ordering::Greater => self.front = None, } - return; - }; - self.back = Some(to64iter(next_back)); - } + } + return; + }; - if let Some(ref mut back) = self.back { + let mut back = to64iter(next_back); + if back.hi == key { back.advance_back_to(index); } + self.back = Some(back); } } diff --git a/roaring/tests/treemap_iter_advance_to.rs b/roaring/tests/treemap_iter_advance_to.rs index bead176b..0685544f 100644 --- a/roaring/tests/treemap_iter_advance_to.rs +++ b/roaring/tests/treemap_iter_advance_to.rs @@ -33,6 +33,31 @@ fn to_next_bitmap() { assert_eq!(i.next(), None); } +#[test] +fn advance_to_later_bitmap_after_iteration_started() { + let later_bitmap = 1u64 << 32; + let bm = RoaringTreemap::from([1, 3, later_bitmap]); + let mut i = bm.iter(); + + assert_eq!(i.next(), Some(1)); + i.advance_to(later_bitmap); + + assert_eq!(i.next(), Some(later_bitmap)); + assert_eq!(i.next(), None); +} + +#[test] +fn advance_to_missing_bitmap() { + let later_bitmap = 1u64 << 32; + let bm = RoaringTreemap::from([later_bitmap]); + let mut i = bm.iter(); + + i.advance_to(3); + + assert_eq!(i.next(), Some(later_bitmap)); + assert_eq!(i.next(), None); +} + #[test] fn iter_back_basic() { let bm = RoaringTreemap::from([1, 2, 3, 4, 11, 12, 13, 14]); @@ -47,6 +72,31 @@ fn iter_back_basic() { assert_eq!(i.next_back(), None); } +#[test] +fn advance_back_to_earlier_bitmap_after_iteration_started() { + let later_bitmap = 1u64 << 32; + let bm = RoaringTreemap::from([1, later_bitmap, later_bitmap + 2]); + let mut i = bm.iter(); + + assert_eq!(i.next_back(), Some(later_bitmap + 2)); + i.advance_back_to(1); + + assert_eq!(i.next_back(), Some(1)); + assert_eq!(i.next_back(), None); +} + +#[test] +fn advance_back_to_missing_bitmap() { + let later_bitmap = 1u64 << 32; + let bm = RoaringTreemap::from([3]); + let mut i = bm.iter(); + + i.advance_back_to(later_bitmap); + + assert_eq!(i.next_back(), Some(3)); + assert_eq!(i.next_back(), None); +} + #[test] fn iter_advance_past_end() { let bm = RoaringTreemap::from([1, 2, 3, 4, 11, 12, 13, 14]);