Merge #303

303: Introduce peek_nth r=jswrenn a=davidcornu

I've been implementing Bob Nystrom's [Crafting Interpreters](https://www.craftinginterpreters.com) in Rust as a way to learn the language and stumbled upon [`multipeek`](https://docs.rs/itertools/0.7.8/itertools/fn.multipeek.html) while looking for a way to `peek` more than one value ahead (as is required in order to [tokenize number literals](http://www.craftinginterpreters.com/scanning.html#number-literals)).

`multipeek` seemed like a great solution to this but I find it's tricky to use correctly. Because it makes `peek` stateful, you have to be careful to call [`reset_peek`](https://docs.rs/itertools/0.7.8/itertools/structs/struct.MultiPeek.html#method.reset_peek) whenever handing control of the iterator to another part of the code or risk running into bugs when `peek` doesn't return what you expect.

To work around this I implemented `peek_nth` which allows you to `peek` more than one iteration into the future without relying on a cursor. I have a working version in my project but I figured I'd submit it here for discussion.

**Example usage**

```rust
let xs = vec![1,2,3];
let mut iter = peek_nth(xs.iter());

assert_eq!(iter.peek(), 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);
```

Note - `peek` equivalent to `peek_nth(0)` and is included for convenience.

----

⚠️ I'm _very_ new to Rust and took a lot of inspiration from`multipeek` to implement this. I'm more than willing to rewrite/fix any mistakes I may have made 😅.

Co-authored-by: David Cornu <[email protected]>

Author: bors[bot]

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,102 @@
+use crate::size_hint;
+use crate::PeekingNext;
+use std::collections::VecDeque;
+use std::iter::Fuse;
+
+/// 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> {
+        let unbuffered_items = (n + 1).saturating_sub(self.buf.len());
+
+        self.buf.extend(self.iter.by_ref().take(unbuffered_items));
+
+        self.buf.get(n)
+    }
+}
+
+impl<I> Iterator for PeekNth<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<I::Item> {
+        self.buf.pop_front().or_else(|| self.iter.next())
+    }
+
+    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,
+    {
+        self.peek().filter(|item| accept(item))?;
+        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);