Add Peekable::peek_mut

3 files changed, 49 insertions, 0 deletions

diff --git a/library/core/src/iter/adapters/peekable.rs b/library/core/src/iter/adapters/peekable.rs
index e7fb3abc942..ebdc2555db2 100644
--- a/library/core/src/iter/adapters/peekable.rs
+++ b/library/core/src/iter/adapters/peekable.rs
@@ -216,6 +216,43 @@ impl<I: Iterator> Peekable<I> {
         self.peeked.get_or_insert_with(|| iter.next()).as_ref()
     }
 
+    /// Returns a mutable reference to the next() value without advancing the iterator.
+    ///
+    /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
+    /// But if the iteration is over, `None` is returned.
+    ///
+    /// Because `peek_mut()` returns a reference, and many iterators iterate over
+    /// references, there can be a possibly confusing situation where the
+    /// return value is a double reference. You can see this effect in the examples
+    /// below.
+    ///
+    /// [`next`]: Iterator::next
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(peekable_peek_mut)]
+    /// let mut iter = [1, 2, 3].iter().peekable();
+    ///
+    /// assert_eq!(iter.peek_mut(), Some(&mut &1));
+    /// assert_eq!(iter.next(), Some(&1));
+    ///
+    /// // Peek into the iterator and modify the value which will be returned next
+    /// if let Some(mut p) = iter.peek_mut() {
+    ///     if *p == &2 {
+    ///         *p = &5;
+    ///     }
+    /// }
+    ///
+    /// assert_eq!(iter.collect::<Vec<_>>(), vec![&5, &3]);
+    /// ```
+    #[inline]
+    #[unstable(feature = "peekable_peek_mut", issue = "78302")]
+    pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
+        let iter = &mut self.iter;
+        self.peeked.get_or_insert_with(|| iter.next()).as_mut()
+    }
+
     /// Consume and return the next value of this iterator if a condition is true.
     ///
     /// If `func` returns `true` for the next value of this iterator, consume and return it.
diff --git a/library/core/tests/iter.rs b/library/core/tests/iter.rs
index 75ca897cadc..6b8a989fa42 100644
--- a/library/core/tests/iter.rs
+++ b/library/core/tests/iter.rs
@@ -1134,6 +1134,17 @@ fn test_iterator_peekable_next_if_eq() {
     assert_eq!(it.next_if_eq(""), None);
 }
 
+#[test]
+fn test_iterator_peekable_mut() {
+    let mut it = vec![1, 2, 3].into_iter().peekable();
+    if let Some(p) = it.peek_mut() {
+        if *p == 1 {
+            *p = 5;
+        }
+    }
+    assert_eq!(it.collect::<Vec<_>>(), vec![5, 2, 3]);
+}
+
 /// This is an iterator that follows the Iterator contract,
 /// but it is not fused. After having returned None once, it will start
 /// producing elements if .next() is called again.
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index 14ef03fd53e..1efb3b74118 100644
--- a/library/core/tests/lib.rs
+++ b/library/core/tests/lib.rs
@@ -56,6 +56,7 @@
 #![feature(unwrap_infallible)]
 #![feature(option_unwrap_none)]
 #![feature(peekable_next_if)]
+#![feature(peekable_peek_mut)]
 #![feature(partition_point)]
 #![feature(once_cell)]
 #![feature(unsafe_block_in_unsafe_fn)]