From 133d6983a1fe3f7639d780049425b1ca3db7af00 Mon Sep 17 00:00:00 2001 From: David Cornu Date: Sun, 2 Sep 2018 20:57:51 -0400 Subject: [PATCH] Introduce peek_nth --- src/free.rs | 2 + src/lib.rs | 4 ++ src/peek_nth.rs | 125 ++++++++++++++++++++++++++++++++++++++++++++++ tests/quick.rs | 10 ++++ tests/test_std.rs | 65 ++++++++++++++++++++++++ 5 files changed, 206 insertions(+) create mode 100644 src/peek_nth.rs diff --git a/src/free.rs b/src/free.rs index 1563679f9..a9cd3d9c5 100644 --- a/src/free.rs +++ b/src/free.rs @@ -19,6 +19,8 @@ pub use put_back_n_impl::put_back_n; #[cfg(feature = "use_std")] pub use multipeek_impl::multipeek; #[cfg(feature = "use_std")] +pub use peek_nth::peek_nth; +#[cfg(feature = "use_std")] pub use kmerge_impl::kmerge; pub use zip_eq_impl::zip_eq; pub use merge_join::merge_join_by; diff --git a/src/lib.rs b/src/lib.rs index c6518ff60..a81c32718 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -83,6 +83,8 @@ pub mod structs { pub use merge_join::MergeJoinBy; #[cfg(feature = "use_std")] pub use multipeek_impl::MultiPeek; + #[cfg(feature = "use_std")] + pub use peek_nth::PeekNth; pub use pad_tail::PadUsing; pub use peeking_take_while::PeekingTakeWhile; pub use process_results_impl::ProcessResults; @@ -142,6 +144,8 @@ mod minmax; #[cfg(feature = "use_std")] mod multipeek_impl; mod pad_tail; +#[cfg(feature = "use_std")] +mod peek_nth; mod peeking_take_while; mod process_results_impl; #[cfg(feature = "use_std")] diff --git a/src/peek_nth.rs b/src/peek_nth.rs new file mode 100644 index 000000000..e86770838 --- /dev/null +++ b/src/peek_nth.rs @@ -0,0 +1,125 @@ +use size_hint; +use std::collections::VecDeque; +use std::iter::Fuse; +use PeekingNext; + +/// See [`peek_nth()`](../fn.peek_nth.html) for more information. +#[derive(Clone, Debug)] +pub struct PeekNth +where + I: Iterator, +{ + iter: Fuse, + buf: VecDeque, +} + +/// A drop-in replacement for `std::iter::Peekable` which adds a `peek_nth` +/// method allowing the user to `peek` at a value several iterations forward +/// without advancing the base iterator. +/// +/// This differs from `multipeek` in that subsequent calls to `peek` or +/// `peek_nth` will always return the same value until `next` is called +/// (making `reset_peek` unnecessary). +pub fn peek_nth(iterable: I) -> PeekNth +where + I: IntoIterator, +{ + PeekNth { + iter: iterable.into_iter().fuse(), + buf: VecDeque::new(), + } +} + +impl PeekNth +where + I: Iterator, +{ + /// Works exactly like the `peek` method in `std::iter::Peekable` + pub fn peek(&mut self) -> Option<&I::Item> { + self.peek_nth(0) + } + + /// Returns a reference to the `nth` value without advancing the iterator. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ```rust + /// use itertools::peek_nth; + /// + /// let xs = vec![1,2,3]; + /// let mut iter = peek_nth(xs.iter()); + /// + /// assert_eq!(iter.peek_nth(0), Some(&&1)); + /// assert_eq!(iter.next(), Some(&1)); + /// + /// // The iterator does not advance even if we call `peek_nth` multiple times + /// assert_eq!(iter.peek_nth(0), Some(&&2)); + /// assert_eq!(iter.peek_nth(1), Some(&&3)); + /// assert_eq!(iter.next(), Some(&2)); + /// + /// // Calling `peek_nth` past the end of the iterator will return `None` + /// assert_eq!(iter.peek_nth(1), None); + /// ``` + pub fn peek_nth(&mut self, n: usize) -> Option<&I::Item> { + if n < self.buf.len() { + return Some(&self.buf[n]); + } + + for _ in self.buf.len()..=n { + match self.iter.next() { + Some(next_item) => { + self.buf.push_back(next_item); + } + None => return None, + } + } + + Some(&self.buf[n]) + } +} + +impl Iterator for PeekNth +where + I: Iterator, +{ + type Item = I::Item; + + fn next(&mut self) -> Option { + if self.buf.is_empty() { + self.iter.next() + } else { + self.buf.pop_front() + } + } + + fn size_hint(&self) -> (usize, Option) { + size_hint::add_scalar(self.iter.size_hint(), self.buf.len()) + } +} + +impl ExactSizeIterator for PeekNth where I: ExactSizeIterator {} + +impl PeekingNext for PeekNth +where + I: Iterator, +{ + fn peeking_next(&mut self, accept: F) -> Option + where + F: FnOnce(&Self::Item) -> bool, + { + if self.buf.is_empty() { + if let Some(r) = self.peek() { + if !accept(r) { + return None; + } + } + } else if let Some(r) = self.buf.get(0) { + if !accept(r) { + return None; + } + } + self.next() + } +} diff --git a/tests/quick.rs b/tests/quick.rs index 685975807..51eda51fd 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -22,6 +22,7 @@ use itertools::free::{ cloned, enumerate, multipeek, + peek_nth, put_back, put_back_n, rciter, @@ -446,6 +447,15 @@ quickcheck! { exact_size(it) } + fn size_peek_nth(a: Iter, s: u8) -> bool { + let mut it = peek_nth(a); + // peek a few times + for n in 0..s { + it.peek_nth(n as usize); + } + exact_size(it) + } + fn equal_merge(a: Vec, b: Vec) -> bool { let mut sa = a.clone(); let mut sb = b.clone(); diff --git a/tests/test_std.rs b/tests/test_std.rs index 35e0529fa..7ad34660e 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -6,6 +6,7 @@ use it::flatten; use it::Itertools; use it::multizip; use it::multipeek; +use it::peek_nth; use it::free::rciter; use it::free::put_back_n; use it::FoldWhile; @@ -351,6 +352,70 @@ fn test_multipeek_peeking_next() { assert_eq!(mp.peek(), None); } +#[test] +fn test_peek_nth() { + let nums = vec![1u8,2,3,4,5]; + + let iter = peek_nth(nums.iter().map(|&x| x)); + assert_eq!(nums, iter.collect::>()); + + let mut iter = peek_nth(nums.iter().map(|&x| x)); + + assert_eq!(iter.peek_nth(0), Some(&1)); + assert_eq!(iter.peek_nth(0), Some(&1)); + assert_eq!(iter.next(), Some(1)); + + assert_eq!(iter.peek_nth(0), Some(&2)); + assert_eq!(iter.peek_nth(1), Some(&3)); + assert_eq!(iter.next(), Some(2)); + + assert_eq!(iter.peek_nth(0), Some(&3)); + assert_eq!(iter.peek_nth(1), Some(&4)); + assert_eq!(iter.peek_nth(2), Some(&5)); + assert_eq!(iter.peek_nth(3), None); + + assert_eq!(iter.next(), Some(3)); + assert_eq!(iter.next(), Some(4)); + + assert_eq!(iter.peek_nth(0), Some(&5)); + assert_eq!(iter.peek_nth(1), None); + assert_eq!(iter.next(), Some(5)); + assert_eq!(iter.next(), None); + + assert_eq!(iter.peek_nth(0), None); + assert_eq!(iter.peek_nth(1), None); +} + +#[test] +fn test_peek_nth_peeking_next() { + use it::PeekingNext; + let nums = vec![1u8,2,3,4,5,6,7]; + let mut iter = peek_nth(nums.iter().map(|&x| x)); + + assert_eq!(iter.peeking_next(|&x| x != 0), Some(1)); + assert_eq!(iter.next(), Some(2)); + + assert_eq!(iter.peek_nth(0), Some(&3)); + assert_eq!(iter.peek_nth(1), Some(&4)); + assert_eq!(iter.peeking_next(|&x| x == 3), Some(3)); + assert_eq!(iter.peek(), Some(&4)); + + assert_eq!(iter.peeking_next(|&x| x != 4), None); + assert_eq!(iter.peeking_next(|&x| x == 4), Some(4)); + assert_eq!(iter.peek_nth(0), Some(&5)); + assert_eq!(iter.peek_nth(1), Some(&6)); + + assert_eq!(iter.peeking_next(|&x| x != 5), None); + assert_eq!(iter.peek(), Some(&5)); + + assert_eq!(iter.peeking_next(|&x| x == 5), Some(5)); + assert_eq!(iter.peeking_next(|&x| x == 6), Some(6)); + assert_eq!(iter.peek_nth(0), Some(&7)); + assert_eq!(iter.peek_nth(1), None); + assert_eq!(iter.next(), Some(7)); + assert_eq!(iter.peek(), None); +} + #[test] fn pad_using() { it::assert_equal((0..0).pad_using(1, |_| 1), 1..2);