From 55f5f8517d2d5d8e3b02653f10df10971567bc1d Mon Sep 17 00:00:00 2001 From: Zachary Dremann Date: Tue, 15 Oct 2024 23:25:00 -0400 Subject: [PATCH] add some more useful interator overrides This adds more efficent overrides for `count`, `nth` and `nth_back`, and implements `FusedIterator` --- roaring/src/bitmap/iter.rs | 131 +++++++++++++++++++++++++++++++++++++ roaring/tests/iter.rs | 47 +++++++++++++ 2 files changed, 178 insertions(+) diff --git a/roaring/src/bitmap/iter.rs b/roaring/src/bitmap/iter.rs index 1593ab40..70ef27fa 100644 --- a/roaring/src/bitmap/iter.rs +++ b/roaring/src/bitmap/iter.rs @@ -1,4 +1,5 @@ use alloc::vec; +use core::iter::FusedIterator; use core::slice; use super::container::{self, Container}; @@ -96,6 +97,43 @@ impl Iterator for Iter<'_> { }; init } + + fn count(self) -> usize + where + Self: Sized, + { + let mut count = self.front.map_or(0, Iterator::count); + count += self.containers.map(|container| container.len() as usize).sum::(); + count += self.back.map_or(0, Iterator::count); + count + } + + fn nth(&mut self, n: usize) -> Option { + let mut n = n; + let nth_advance = |it: &mut container::Iter| { + let len = it.len(); + if n < len { + it.nth(n) + } else { + n -= len; + None + } + }; + if let Some(x) = and_then_or_clear(&mut self.front, nth_advance) { + return Some(x); + } + for container in self.containers.by_ref() { + let len = container.len() as usize; + if n < len { + let mut front_iter = container.into_iter(); + let result = front_iter.nth(n); + self.front = Some(front_iter); + return result; + } + n -= len; + } + and_then_or_clear(&mut self.back, |it| it.nth(n)) + } } impl DoubleEndedIterator for Iter<'_> { @@ -128,10 +166,38 @@ impl DoubleEndedIterator for Iter<'_> { }; init } + + fn nth_back(&mut self, n: usize) -> Option { + let mut n = n; + let nth_advance = |it: &mut container::Iter| { + let len = it.len(); + if n < len { + it.nth_back(n) + } else { + n -= len; + None + } + }; + if let Some(x) = and_then_or_clear(&mut self.back, nth_advance) { + return Some(x); + } + for container in self.containers.by_ref().rev() { + let len = container.len() as usize; + if n < len { + let mut front_iter = container.into_iter(); + let result = front_iter.nth_back(n); + self.back = Some(front_iter); + return result; + } + n -= len; + } + and_then_or_clear(&mut self.front, |it| it.nth_back(n)) + } } #[cfg(target_pointer_width = "64")] impl ExactSizeIterator for Iter<'_> {} +impl FusedIterator for Iter<'_> {} impl Iterator for IntoIter { type Item = u32; @@ -170,6 +236,43 @@ impl Iterator for IntoIter { }; init } + + fn count(self) -> usize + where + Self: Sized, + { + let mut count = self.front.map_or(0, Iterator::count); + count += self.containers.map(|container| container.len() as usize).sum::(); + count += self.back.map_or(0, Iterator::count); + count + } + + fn nth(&mut self, n: usize) -> Option { + let mut n = n; + let nth_advance = |it: &mut container::Iter| { + let len = it.len(); + if n < len { + it.nth(n) + } else { + n -= len; + None + } + }; + if let Some(x) = and_then_or_clear(&mut self.front, nth_advance) { + return Some(x); + } + for container in self.containers.by_ref() { + let len = container.len() as usize; + if n < len { + let mut front_iter = container.into_iter(); + let result = front_iter.nth(n); + self.front = Some(front_iter); + return result; + } + n -= len; + } + and_then_or_clear(&mut self.back, |it| it.nth(n)) + } } impl DoubleEndedIterator for IntoIter { @@ -202,10 +305,38 @@ impl DoubleEndedIterator for IntoIter { }; init } + + fn nth_back(&mut self, n: usize) -> Option { + let mut n = n; + let nth_advance = |it: &mut container::Iter| { + let len = it.len(); + if n < len { + it.nth_back(n) + } else { + n -= len; + None + } + }; + if let Some(x) = and_then_or_clear(&mut self.back, nth_advance) { + return Some(x); + } + for container in self.containers.by_ref().rev() { + let len = container.len() as usize; + if n < len { + let mut front_iter = container.into_iter(); + let result = front_iter.nth_back(n); + self.back = Some(front_iter); + return result; + } + n -= len; + } + and_then_or_clear(&mut self.front, |it| it.nth_back(n)) + } } #[cfg(target_pointer_width = "64")] impl ExactSizeIterator for IntoIter {} +impl FusedIterator for IntoIter {} impl RoaringBitmap { /// Iterator over each value stored in the RoaringBitmap, guarantees values are ordered by value. diff --git a/roaring/tests/iter.rs b/roaring/tests/iter.rs index 86a83245..05591681 100644 --- a/roaring/tests/iter.rs +++ b/roaring/tests/iter.rs @@ -81,6 +81,53 @@ proptest! { } } +proptest! { + #[test] + fn nth(values in btree_set(any::(), ..=10_000), nth in 0..10_005usize) { + let bitmap = RoaringBitmap::from_sorted_iter(values.iter().cloned()).unwrap(); + let mut orig_iter = bitmap.iter().fuse(); + let mut iter = bitmap.iter(); + + for _ in 0..nth { + if orig_iter.next().is_none() { + break; + } + } + let expected = orig_iter.next(); + assert_eq!(expected, iter.nth(nth)); + let expected_next = orig_iter.next(); + assert_eq!(expected_next, iter.next()); + + let mut val_iter = values.into_iter(); + assert_eq!(expected, val_iter.nth(nth)); + assert_eq!(expected_next, val_iter.next()); + } +} + +#[test] +fn huge_nth() { + let bitmap = RoaringBitmap::new(); + let mut iter = bitmap.iter(); + assert_eq!(None, iter.nth(usize::MAX)); +} + +proptest! { + #[test] + fn count(values in btree_set(any::(), ..=10_000), skip in 0..10_005usize) { + let bitmap = RoaringBitmap::from_sorted_iter(values.iter().cloned()).unwrap(); + let mut iter = bitmap.iter(); + + if let Some(n) = skip.checked_sub(1) { + iter.nth(n); + } + let expected_count = values.len().saturating_sub(skip); + let size_hint = iter.size_hint(); + assert_eq!(expected_count, size_hint.0); + assert_eq!(Some(expected_count), size_hint.1); + assert_eq!(expected_count, iter.count()); + } +} + #[test] fn rev_array() { let values = 0..100;