Repository: rust-itertools/itertools (language: rust)

Issue:
Support for `next_if` and `next_if_eq` on PeekNth
Currently the documentation for [`itertools::peek_nth`](https://docs.rs/itertools/0.11.0/itertools/fn.peek_nth.html) says:

> A drop-in replacement for [std::iter::Peekable](https://doc.rust-lang.org/nightly/core/iter/adapters/peekable/struct.Peekable.html) which adds a peek_nth method allowing the user to peek at a value several iterations forward without advancing the base iterator.

It is however not a drop-in replacement at the moment, since the methods [`next_if`](https://doc.rust-lang.org/nightly/core/iter/struct.Peekable.html#method.next_if) and [`next_if_eq`](https://doc.rust-lang.org/nightly/core/iter/struct.Peekable.html#method.next_if_eq) isn't implemented.


Discussion / hints:

Fix patch:
diff --git a/src/peek_nth.rs b/src/peek_nth.rs
index 1ae428ba0..5f7fb9396 100644
--- a/src/peek_nth.rs
+++ b/src/peek_nth.rs
@@ -115,6 +115,27 @@ where
 
         self.buf.get_mut(n)
     }
+
+    /// Works exactly like the `next_if` method in `std::iter::Peekable`
+    pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> {
+        match self.next() {
+            Some(item) if func(&item) => Some(item),
+            Some(item) => {
+                self.buf.push_front(item);
+                None
+            }
+            _ => None,
+        }
+    }
+
+    /// Works exactly like the `next_if_eq` method in `std::iter::Peekable`
+    pub fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item>
+    where
+        T: ?Sized,
+        I::Item: PartialEq<T>,
+    {
+        self.next_if(|next| next == expected)
+    }
 }
 
 impl<I> Iterator for PeekNth<I>


Test patch:
diff --git a/tests/quick.rs b/tests/quick.rs
index 0d7cd9307..f0c8f2d9e 100644
--- a/tests/quick.rs
+++ b/tests/quick.rs
@@ -782,6 +782,30 @@ quickcheck! {
         assert_eq!(it.next(), None);
         assert_eq!(it.next(), None);
     }
+
+    fn peek_nth_next_if(a: Vec<u8>) -> () {
+        let mut it = peek_nth(a.clone());
+        for (idx, mut value) in a.iter().copied().enumerate() {
+            let should_be_none = it.next_if(|x| x != &value);
+            assert_eq!(should_be_none, None);
+            if value % 5 == 0 {
+                // Sometimes, peek up to 3 further.
+                let n = value as usize % 3;
+                let nth = it.peek_nth(n);
+                assert_eq!(nth, a.get(idx + n));
+            } else if value % 5 == 1 {
+                // Sometimes, peek next element mutably.
+                if let Some(v) = it.peek_mut() {
+                    *v = v.wrapping_sub(1);
+                    let should_be_none = it.next_if_eq(&value);
+                    assert_eq!(should_be_none, None);
+                    value = value.wrapping_sub(1);
+                }
+            }
+            let eq = it.next_if_eq(&value);
+            assert_eq!(eq, Some(value));
+        }
+    }
 }
 
 quickcheck! {
diff --git a/tests/test_std.rs b/tests/test_std.rs
index 3ef90f39d..732790ef4 100644
--- a/tests/test_std.rs
+++ b/tests/test_std.rs
@@ -738,6 +738,35 @@ fn test_peek_nth_peeking_next() {
     assert_eq!(iter.peek(), None);
 }
 
+#[test]
+fn test_peek_nth_next_if() {
+    let nums = vec![1u8, 2, 3, 4, 5, 6, 7];
+    let mut iter = peek_nth(nums.iter().copied());
+
+    assert_eq!(iter.next_if(|&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.next_if_eq(&3), Some(3));
+    assert_eq!(iter.peek(), Some(&4));
+
+    assert_eq!(iter.next_if(|&x| x != 4), None);
+    assert_eq!(iter.next_if_eq(&4), Some(4));
+    assert_eq!(iter.peek_nth(0), Some(&5));
+    assert_eq!(iter.peek_nth(1), Some(&6));
+
+    assert_eq!(iter.next_if(|&x| x != 5), None);
+    assert_eq!(iter.peek(), Some(&5));
+
+    assert_eq!(iter.next_if(|&x| x % 2 == 1), Some(5));
+    assert_eq!(iter.next_if_eq(&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);