1 files changed, 16 insertions, 10 deletions
diff --git a/library/alloc/src/vec/peek_mut.rs b/library/alloc/src/vec/peek_mut.rs
index c0dd941ed39..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,42 +12,47 @@ 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 }) }
     }
 
     /// Removes the peeked value from the vector and returns it.
     #[unstable(feature = "vec_peek_mut", issue = "122742")]
-    pub fn pop(self) -> T {
+    pub fn pop(this: Self) -> T {
         // SAFETY: PeekMut is only constructed if the vec is non-empty
-        unsafe { self.vec.pop().unwrap_unchecked() }
+        unsafe { this.vec.pop().unwrap_unchecked() }
     }
 }
 
 #[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