Plumb Allocator generic into `std::vec::PeekMut`

3 files changed, 51 insertions, 44 deletions

diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index 2e40227a058..1b80b43abc3 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -760,33 +760,6 @@ impl<T> Vec<T> {
         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)
-    }
-
     /// Decomposes a `Vec<T>` into its raw components: `(pointer, length, capacity)`.
     ///
     /// Returns the raw pointer to the underlying data, the length of
@@ -2747,6 +2720,33 @@ impl<T, A: Allocator> Vec<T, A> {
         if predicate(last) { self.pop() } else { None }
     }
 
+    /// 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, A>> {
+        PeekMut::new(self)
+    }
+
     /// Moves all the elements of `other` into `self`, leaving `other` empty.
     ///
     /// # Panics
diff --git a/library/alloc/src/vec/peek_mut.rs b/library/alloc/src/vec/peek_mut.rs
index caeaf2799d7..979bcaa1111 100644
--- a/library/alloc/src/vec/peek_mut.rs
+++ b/library/alloc/src/vec/peek_mut.rs
@@ -1,6 +1,7 @@
 use core::ops::{Deref, DerefMut};
 
 use super::Vec;
+use crate::alloc::{Allocator, Global};
 use crate::fmt;
 
 /// Structure wrapping a mutable reference to the last item in a
@@ -11,19 +12,23 @@ use crate::fmt;
 ///
 /// [`peek_mut`]: Vec::peek_mut
 #[unstable(feature = "vec_peek_mut", issue = "122742")]
-pub struct PeekMut<'a, T> {
-    vec: &'a mut Vec<T>,
+pub struct PeekMut<
+    'a,
+    T,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
+> {
+    vec: &'a mut Vec<T, A>,
 }
 
 #[unstable(feature = "vec_peek_mut", issue = "122742")]
-impl<T: fmt::Debug> fmt::Debug for PeekMut<'_, T> {
+impl<T: fmt::Debug, A: Allocator> fmt::Debug for PeekMut<'_, T, A> {
     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> {
+impl<'a, T, A: Allocator> PeekMut<'a, T, A> {
+    pub(super) fn new(vec: &'a mut Vec<T, A>) -> Option<Self> {
         if vec.is_empty() { None } else { Some(Self { vec }) }
     }
 
@@ -36,17 +41,18 @@ impl<'a, T> PeekMut<'a, T> {
 }
 
 #[unstable(feature = "vec_peek_mut", issue = "122742")]
-impl<'a, T> Deref for PeekMut<'a, T> {
+impl<'a, T, A: Allocator> Deref for PeekMut<'a, T, A> {
     type Target = T;
 
     fn deref(&self) -> &Self::Target {
+        let idx = self.vec.len() - 1;
         // SAFETY: PeekMut is only constructed if the vec is non-empty
-        unsafe { self.vec.get_unchecked(self.vec.len() - 1) }
+        unsafe { self.vec.get_unchecked(idx) }
     }
 }
 
 #[unstable(feature = "vec_peek_mut", issue = "122742")]
-impl<'a, T> DerefMut for PeekMut<'a, T> {
+impl<'a, T, A: Allocator> DerefMut for PeekMut<'a, T, A> {
     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
diff --git a/library/alloctests/tests/vec.rs b/library/alloctests/tests/vec.rs
index 404eb49e1ea..33a34daccbf 100644
--- a/library/alloctests/tests/vec.rs
+++ b/library/alloctests/tests/vec.rs
@@ -2643,15 +2643,16 @@ fn test_peek_mut() {
     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);
-        PeekMut::pop(p);
-        assert_eq!(vec.len(), 1);
-    } else {
-        unreachable!()
-    }
+    let mut p = vec.peek_mut().unwrap();
+    assert_eq!(*p, 2);
+    *p = 0;
+    assert_eq!(*p, 0);
+    drop(p);
+    assert_eq!(vec, vec![1, 0]);
+    let p = vec.peek_mut().unwrap();
+    let p = PeekMut::pop(p);
+    assert_eq!(p, 0);
+    assert_eq!(vec, vec![1]);
 }
 
 /// This assortment of tests, in combination with miri, verifies we handle UB on fishy arguments