Add peek_mut and peek_nth_mut to PeekNth

src/peek_nth.rs

@@ -39,25 +39,30 @@ where
         self.peek_nth(0)
     }
 
+    /// Works exactly like the `peek_mut` method in `std::iter::Peekable`
+    pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
+        self.peek_nth_mut(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());
+    /// let xs = vec![1, 2, 3];
+    /// let mut iter = peek_nth(xs.into_iter());
     ///
-    /// assert_eq!(iter.peek_nth(0), Some(&&1));
-    /// assert_eq!(iter.next(), Some(&1));
+    /// 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));
+    /// 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);
@@ -69,6 +74,47 @@ where
 
         self.buf.get(n)
     }
+
+    /// Returns a mutable reference to the `nth` value without advancing the iterator.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// use itertools::peek_nth;
+    ///
+    /// let xs = vec![1, 2, 3, 4, 5];
+    /// let mut iter = peek_nth(xs.into_iter());
+    ///
+    /// assert_eq!(iter.peek_nth_mut(0), Some(&mut 1));
+    /// assert_eq!(iter.next(), Some(1));
+    ///
+    /// // The iterator does not advance even if we call `peek_nth_mut` multiple times
+    /// assert_eq!(iter.peek_nth_mut(0), Some(&mut 2));
+    /// assert_eq!(iter.peek_nth_mut(1), Some(&mut 3));
+    /// assert_eq!(iter.next(), Some(2));
+    ///
+    /// // Peek into the iterator and set the value behind the mutable reference.
+    /// if let Some(p) = iter.peek_nth_mut(1) {
+    ///     assert_eq!(*p, 4);
+    ///     *p = 9;
+    /// }
+    ///
+    /// // The value we put in reappears as the iterator continues.
+    /// assert_eq!(iter.next(), Some(3));
+    /// assert_eq!(iter.next(), Some(9));
+    ///
+    /// // Calling `peek_nth_mut` past the end of the iterator will return `None`
+    /// assert_eq!(iter.peek_nth_mut(1), None);
+    /// ```
+    pub fn peek_nth_mut(&mut self, n: usize) -> Option<&mut 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_mut(n)
+    }
 }
 
 impl<I> Iterator for PeekNth<I>

tests/quick.rs

@@ -758,6 +758,32 @@ quickcheck! {
     }
 }
 
+quickcheck! {
+    fn correct_peek_nth(mut a: Vec<u16>) -> () {
+        let mut it = peek_nth(a.clone());
+        for start_pos in 0..a.len() + 2 {
+            for real_idx in start_pos..a.len() + 2 {
+                let peek_idx = real_idx - start_pos;
+                assert_eq!(it.peek_nth(peek_idx), a.get(real_idx));
+                assert_eq!(it.peek_nth_mut(peek_idx), a.get_mut(real_idx));
+            }
+            assert_eq!(it.next(), a.get(start_pos).copied());
+        }
+    }
+
+    fn peek_nth_mut_replace(a: Vec<u16>, b: Vec<u16>) -> () {
+        let mut it = peek_nth(a.iter());
+        for i in 0..a.len().min(b.len()) {
+            *it.peek_nth_mut(i).unwrap() = &b[i];
+        }
+        for i in 0..a.len() {
+            assert_eq!(it.next().unwrap(), b.get(i).unwrap_or(&a[i]));
+        }
+        assert_eq!(it.next(), None);
+        assert_eq!(it.next(), None);
+    }
+}
+
 quickcheck! {
     fn dedup_via_coalesce(a: Vec<i32>) -> bool {
         let mut b = a.clone();