Introduce peek_nth

src/free.rs

@@ -22,6 +22,8 @@ pub use crate::put_back_n_impl::put_back_n;
 #[cfg(feature = "use_std")]
 pub use crate::multipeek_impl::multipeek;
 #[cfg(feature = "use_std")]
+pub use crate::peek_nth::peek_nth;
+#[cfg(feature = "use_std")]
 pub use crate::kmerge_impl::kmerge;
 pub use crate::zip_eq_impl::zip_eq;
 pub use crate::merge_join::merge_join_by;

src/lib.rs

@@ -113,6 +113,8 @@ pub mod structs {
     pub use crate::merge_join::MergeJoinBy;
     #[cfg(feature = "use_std")]
     pub use crate::multipeek_impl::MultiPeek;
+    #[cfg(feature = "use_std")]
+    pub use crate::peek_nth::PeekNth;
     pub use crate::pad_tail::PadUsing;
     pub use crate::peeking_take_while::PeekingTakeWhile;
     #[cfg(feature = "use_std")]
@@ -187,6 +189,8 @@ mod minmax;
 #[cfg(feature = "use_std")]
 mod multipeek_impl;
 mod pad_tail;
+#[cfg(feature = "use_std")]
+mod peek_nth;
 mod peeking_take_while;
 #[cfg(feature = "use_std")]
 mod permutations;

src/peek_nth.rs

@@ -0,0 +1,121 @@
+use crate::size_hint;
+use std::collections::VecDeque;
+use std::iter::Fuse;
+use crate::PeekingNext;
+
+/// See [`peek_nth()`](../fn.peek_nth.html) for more information.
+#[derive(Clone, Debug)]
+pub struct PeekNth<I>
+where
+  I: Iterator,
+{
+  iter: Fuse<I>,
+  buf: VecDeque<I::Item>,
+}
+
+/// 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<I>(iterable: I) -> PeekNth<I::IntoIter>
+where
+  I: IntoIterator,
+{
+  PeekNth {
+    iter: iterable.into_iter().fuse(),
+    buf: VecDeque::new(),
+  }
+}
+
+impl<I> PeekNth<I>
+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> {
+    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<I> Iterator for PeekNth<I>
+where
+  I: Iterator,
+{
+  type Item = I::Item;
+
+  fn next(&mut self) -> Option<I::Item> {
+    if self.buf.is_empty() {
+      self.iter.next()
+    } else {
+      self.buf.pop_front()
+    }
+  }
+
+  fn size_hint(&self) -> (usize, Option<usize>) {
+    size_hint::add_scalar(self.iter.size_hint(), self.buf.len())
+  }
+}
+
+impl<I> ExactSizeIterator for PeekNth<I> where I: ExactSizeIterator {}
+
+impl<I> PeekingNext for PeekNth<I>
+where
+  I: Iterator,
+{
+  fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
+  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()
+  }
+}

tests/quick.rs

@@ -19,6 +19,7 @@ use itertools::free::{
     cloned,
     enumerate,
     multipeek,
+    peek_nth,
     put_back,
     put_back_n,
     rciter,
@@ -487,6 +488,15 @@ quickcheck! {
         exact_size(it)
     }
 
+    fn size_peek_nth(a: Iter<u16, Exact>, 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<i16>, b: Vec<i16>) -> bool {
         let mut sa = a.clone();
         let mut sb = b.clone();

tests/test_std.rs

@@ -3,6 +3,7 @@ use itertools as it;
 use crate::it::Itertools;
 use crate::it::multizip;
 use crate::it::multipeek;
+use crate::it::peek_nth;
 use crate::it::free::rciter;
 use crate::it::free::put_back_n;
 use crate::it::FoldWhile;
@@ -371,6 +372,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::<Vec<_>>());
+
+    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);