Rollup merge of #142046 - Qelxiros:122742-vec_peek_mut, r=cuviper

add Vec::peek_mut

Tracking issue: rust-lang/rust#122742

6 files changed, 107 insertions, 2 deletions

diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index ce7321544b6..5bd82560da7 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -109,6 +109,11 @@ mod in_place_collect;
 
 mod partial_eq;
 
+#[unstable(feature = "vec_peek_mut", issue = "122742")]
+pub use self::peek_mut::PeekMut;
+
+mod peek_mut;
+
 #[cfg(not(no_global_oom_handling))]
 use self::spec_from_elem::SpecFromElem;
 
@@ -729,6 +734,33 @@ impl<T> Vec<T> {
     pub unsafe fn from_parts(ptr: NonNull<T>, length: usize, capacity: usize) -> Self {
         unsafe { Self::from_parts_in(ptr, length, capacity, Global) }
     }
+
+    /// Returns a mutable reference to the last item in the vector, or
+    /// `None` if it is empty.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(vec_peek_mut)]
+    /// let mut vec = Vec::new();
+    /// assert!(vec.peek_mut().is_none());
+    ///
+    /// vec.push(1);
+    /// vec.push(5);
+    /// vec.push(2);
+    /// assert_eq!(vec.last(), Some(&2));
+    /// if let Some(mut val) = vec.peek_mut() {
+    ///     *val = 0;
+    /// }
+    /// assert_eq!(vec.last(), Some(&0));
+    /// ```
+    #[inline]
+    #[unstable(feature = "vec_peek_mut", issue = "122742")]
+    pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T>> {
+        PeekMut::new(self)
+    }
 }
 
 impl<T, A: Allocator> Vec<T, A> {
diff --git a/library/alloc/src/vec/peek_mut.rs b/library/alloc/src/vec/peek_mut.rs
new file mode 100644
index 00000000000..c0dd941ed39
--- /dev/null
+++ b/library/alloc/src/vec/peek_mut.rs
@@ -0,0 +1,55 @@
+use core::ops::{Deref, DerefMut};
+
+use super::Vec;
+use crate::fmt;
+
+/// Structure wrapping a mutable reference to the last item in a
+/// `Vec`.
+///
+/// This `struct` is created by the [`peek_mut`] method on [`Vec`]. See
+/// its documentation for more.
+///
+/// [`peek_mut`]: Vec::peek_mut
+#[unstable(feature = "vec_peek_mut", issue = "122742")]
+pub struct PeekMut<'a, T> {
+    vec: &'a mut Vec<T>,
+}
+
+#[unstable(feature = "vec_peek_mut", issue = "122742")]
+impl<T: fmt::Debug> fmt::Debug for PeekMut<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("PeekMut").field(self.deref()).finish()
+    }
+}
+
+impl<'a, T> PeekMut<'a, T> {
+    pub(crate) fn new(vec: &'a mut Vec<T>) -> Option<Self> {
+        if vec.is_empty() { None } else { Some(Self { vec }) }
+    }
+
+    /// Removes the peeked value from the vector and returns it.
+    #[unstable(feature = "vec_peek_mut", issue = "122742")]
+    pub fn pop(self) -> T {
+        // SAFETY: PeekMut is only constructed if the vec is non-empty
+        unsafe { self.vec.pop().unwrap_unchecked() }
+    }
+}
+
+#[unstable(feature = "vec_peek_mut", issue = "122742")]
+impl<'a, T> Deref for PeekMut<'a, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: PeekMut is only constructed if the vec is non-empty
+        unsafe { self.vec.get_unchecked(self.vec.len() - 1) }
+    }
+}
+
+#[unstable(feature = "vec_peek_mut", issue = "122742")]
+impl<'a, T> DerefMut for PeekMut<'a, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        let idx = self.vec.len() - 1;
+        // SAFETY: PeekMut is only constructed if the vec is non-empty
+        unsafe { self.vec.get_unchecked_mut(idx) }
+    }
+}
diff --git a/library/alloctests/tests/lib.rs b/library/alloctests/tests/lib.rs
index 38309585fad..a41162ecd51 100644
--- a/library/alloctests/tests/lib.rs
+++ b/library/alloctests/tests/lib.rs
@@ -40,6 +40,7 @@
 #![feature(vec_deque_truncate_front)]
 #![feature(unique_rc_arc)]
 #![feature(macro_metavar_expr_concat)]
+#![feature(vec_peek_mut)]
 #![allow(internal_features)]
 #![deny(fuzzy_provenance_casts)]
 #![deny(unsafe_op_in_unsafe_fn)]
diff --git a/library/alloctests/tests/vec.rs b/library/alloctests/tests/vec.rs
index f430d979fa8..51b49b8edb3 100644
--- a/library/alloctests/tests/vec.rs
+++ b/library/alloctests/tests/vec.rs
@@ -2698,6 +2698,23 @@ fn test_pop_if_mutates() {
     assert_eq!(v, [2]);
 }
 
+#[test]
+fn test_peek_mut() {
+    let mut vec = Vec::new();
+    assert!(vec.peek_mut().is_none());
+    vec.push(1);
+    vec.push(2);
+    if let Some(mut p) = vec.peek_mut() {
+        assert_eq!(*p, 2);
+        *p = 0;
+        assert_eq!(*p, 0);
+        p.pop();
+        assert_eq!(vec.len(), 1);
+    } else {
+        unreachable!()
+    }
+}
+
 /// This assortment of tests, in combination with miri, verifies we handle UB on fishy arguments
 /// in the stdlib. Draining and extending the allocation are fairly well-tested earlier, but
 /// `vec.insert(usize::MAX, val)` once slipped by!
diff --git a/tests/ui/borrowck/issue-47646.stderr b/tests/ui/borrowck/issue-47646.stderr
index 85adfc03d10..cfe6f3f3993 100644
--- a/tests/ui/borrowck/issue-47646.stderr
+++ b/tests/ui/borrowck/issue-47646.stderr
@@ -11,7 +11,7 @@ LL |             println!("{:?}", heap);
    |                              ^^^^ immutable borrow occurs here
 ...
 LL |     };
-   |      - ... and the mutable borrow might be used here, when that temporary is dropped and runs the destructor for type `(Option<PeekMut<'_, i32>>, ())`
+   |      - ... and the mutable borrow might be used here, when that temporary is dropped and runs the destructor for type `(Option<std::collections::binary_heap::PeekMut<'_, i32>>, ())`
    |
    = note: this error originates in the macro `$crate::format_args_nl` which comes from the expansion of the macro `println` (in Nightly builds, run with -Z macro-backtrace for more info)
 
diff --git a/tests/ui/borrowck/issue-85581.stderr b/tests/ui/borrowck/issue-85581.stderr
index 80f1f4cb509..5fd457eb8dd 100644
--- a/tests/ui/borrowck/issue-85581.stderr
+++ b/tests/ui/borrowck/issue-85581.stderr
@@ -10,7 +10,7 @@ LL |         Some(_) => { heap.pop(); },
    |                      ^^^^ second mutable borrow occurs here
 ...
 LL | }
-   | - ... and the first borrow might be used here, when that temporary is dropped and runs the destructor for type `Option<PeekMut<'_, i32>>`
+   | - ... and the first borrow might be used here, when that temporary is dropped and runs the destructor for type `Option<std::collections::binary_heap::PeekMut<'_, i32>>`
 
 error: aborting due to 1 previous error