Parametrize BinaryHeap with Allocator

Author: yanchith

library/alloc/src/collections/binary_heap.rs

@@ -143,12 +143,14 @@
 #![allow(missing_docs)]
 #![stable(feature = "rust1", since = "1.0.0")]
 
+use core::alloc::Allocator;
 use core::fmt;
 use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen};
 use core::mem::{self, swap, ManuallyDrop};
 use core::ops::{Deref, DerefMut};
 use core::ptr;
 
+use crate::alloc::Global;
 use crate::collections::TryReserveError;
 use crate::slice;
 use crate::vec::{self, AsVecIntoIter, Vec};
@@ -265,8 +267,11 @@ mod tests;
 /// [peek\_mut]: BinaryHeap::peek_mut
 #[stable(feature = "rust1", since = "1.0.0")]
 #[cfg_attr(not(test), rustc_diagnostic_item = "BinaryHeap")]
-pub struct BinaryHeap<T> {
-    data: Vec<T>,
+pub struct BinaryHeap<
+    T,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
+> {
+    data: Vec<T, A>,
 }
 
 /// Structure wrapping a mutable reference to the greatest item on a
@@ -277,20 +282,24 @@ pub struct BinaryHeap<T> {
 ///
 /// [`peek_mut`]: BinaryHeap::peek_mut
 #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
-pub struct PeekMut<'a, T: 'a + Ord> {
-    heap: &'a mut BinaryHeap<T>,
+pub struct PeekMut<
+    'a,
+    T: 'a + Ord,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + 'a = Global,
+> {
+    heap: &'a mut BinaryHeap<T, A>,
     sift: bool,
 }
 
 #[stable(feature = "collection_debug", since = "1.17.0")]
-impl<T: Ord + fmt::Debug> fmt::Debug for PeekMut<'_, T> {
+impl<T: Ord + 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.heap.data[0]).finish()
     }
 }
 
 #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
-impl<T: Ord> Drop for PeekMut<'_, T> {
+impl<T: Ord, A: Allocator> Drop for PeekMut<'_, T, A> {
     fn drop(&mut self) {
         if self.sift {
             // SAFETY: PeekMut is only instantiated for non-empty heaps.
@@ -300,7 +309,7 @@ impl<T: Ord> Drop for PeekMut<'_, T> {
 }
 
 #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
-impl<T: Ord> Deref for PeekMut<'_, T> {
+impl<T: Ord, A: Allocator> Deref for PeekMut<'_, T, A> {
     type Target = T;
     fn deref(&self) -> &T {
         debug_assert!(!self.heap.is_empty());
@@ -310,7 +319,7 @@ impl<T: Ord> Deref for PeekMut<'_, T> {
 }
 
 #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
-impl<T: Ord> DerefMut for PeekMut<'_, T> {
+impl<T: Ord, A: Allocator> DerefMut for PeekMut<'_, T, A> {
     fn deref_mut(&mut self) -> &mut T {
         debug_assert!(!self.heap.is_empty());
         self.sift = true;
@@ -330,7 +339,7 @@ impl<'a, T: Ord> PeekMut<'a, T> {
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<T: Clone> Clone for BinaryHeap<T> {
+impl<T: Clone, A: Allocator + Clone> Clone for BinaryHeap<T, A> {
     fn clone(&self) -> Self {
         BinaryHeap { data: self.data.clone() }
     }
@@ -350,13 +359,13 @@ impl<T: Ord> Default for BinaryHeap<T> {
 }
 
 #[stable(feature = "binaryheap_debug", since = "1.4.0")]
-impl<T: fmt::Debug> fmt::Debug for BinaryHeap<T> {
+impl<T: fmt::Debug, A: Allocator> fmt::Debug for BinaryHeap<T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_list().entries(self.iter()).finish()
     }
 }
 
-impl<T: Ord> BinaryHeap<T> {
+impl<T: Ord> BinaryHeap<T, Global> {
     /// Creates an empty `BinaryHeap` as a max-heap.
     ///
     /// # Examples
@@ -394,6 +403,52 @@ impl<T: Ord> BinaryHeap<T> {
     pub fn with_capacity(capacity: usize) -> BinaryHeap<T> {
         BinaryHeap { data: Vec::with_capacity(capacity) }
     }
+}
+
+impl<T: Ord, A: Allocator> BinaryHeap<T, A> {
+    /// Creates an empty `BinaryHeap` as a max-heap, using `A` as allocator.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(allocator_api)]
+    ///
+    /// use std::alloc::System;
+    /// use std::collections::BinaryHeap;
+    /// let mut heap = BinaryHeap::new_in(System);
+    /// heap.push(4);
+    /// ```
+    #[unstable(feature = "allocator_api", issue = "32838")]
+    #[must_use]
+    pub fn new_in(alloc: A) -> BinaryHeap<T, A> {
+        BinaryHeap { data: Vec::new_in(alloc) }
+    }
+
+    /// Creates an empty `BinaryHeap` with at least the specified capacity, using `A` as allocator.
+    ///
+    /// The binary heap will be able to hold at least `capacity` elements without
+    /// reallocating. This method is allowed to allocate for more elements than
+    /// `capacity`. If `capacity` is 0, the binary heap will not allocate.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(allocator_api)]
+    ///
+    /// use std::alloc::System;
+    /// use std::collections::BinaryHeap;
+    /// let mut heap = BinaryHeap::with_capacity_in(10, System);
+    /// heap.push(4);
+    /// ```
+    #[unstable(feature = "allocator_api", issue = "32838")]
+    #[must_use]
+    pub fn with_capacity_in(capacity: usize, alloc: A) -> BinaryHeap<T, A> {
+        BinaryHeap { data: Vec::with_capacity_in(capacity, alloc) }
+    }
 
     /// Returns a mutable reference to the greatest item in the binary heap, or
     /// `None` if it is empty.
@@ -425,7 +480,7 @@ impl<T: Ord> BinaryHeap<T> {
     /// If the item is modified then the worst case time complexity is *O*(log(*n*)),
     /// otherwise it's *O*(1).
     #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
-    pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T>> {
+    pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T, A>> {
         if self.is_empty() { None } else { Some(PeekMut { heap: self, sift: false }) }
     }
 
@@ -520,7 +575,7 @@ impl<T: Ord> BinaryHeap<T> {
     /// ```
     #[must_use = "`self` will be dropped if the result is not used"]
     #[stable(feature = "binary_heap_extras_15", since = "1.5.0")]
-    pub fn into_sorted_vec(mut self) -> Vec<T> {
+    pub fn into_sorted_vec(mut self) -> Vec<T, A> {
         let mut end = self.len();
         while end > 1 {
             end -= 1;
@@ -778,7 +833,7 @@ impl<T: Ord> BinaryHeap<T> {
     /// ```
     #[inline]
     #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
-    pub fn drain_sorted(&mut self) -> DrainSorted<'_, T> {
+    pub fn drain_sorted(&mut self) -> DrainSorted<'_, T, A> {
         DrainSorted { inner: self }
     }
 
@@ -821,7 +876,7 @@ impl<T: Ord> BinaryHeap<T> {
     }
 }
 
-impl<T> BinaryHeap<T> {
+impl<T, A: Allocator> BinaryHeap<T, A> {
     /// Returns an iterator visiting all values in the underlying vector, in
     /// arbitrary order.
     ///
@@ -858,7 +913,7 @@ impl<T> BinaryHeap<T> {
     /// assert_eq!(heap.into_iter_sorted().take(2).collect::<Vec<_>>(), [5, 4]);
     /// ```
     #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
-    pub fn into_iter_sorted(self) -> IntoIterSorted<T> {
+    pub fn into_iter_sorted(self) -> IntoIterSorted<T, A> {
         IntoIterSorted { inner: self }
     }
 
@@ -1124,7 +1179,7 @@ impl<T> BinaryHeap<T> {
     /// ```
     #[must_use = "`self` will be dropped if the result is not used"]
     #[stable(feature = "binary_heap_extras_15", since = "1.5.0")]
-    pub fn into_vec(self) -> Vec<T> {
+    pub fn into_vec(self) -> Vec<T, A> {
         self.into()
     }
 
@@ -1195,7 +1250,7 @@ impl<T> BinaryHeap<T> {
     /// ```
     #[inline]
     #[stable(feature = "drain", since = "1.6.0")]
-    pub fn drain(&mut self) -> Drain<'_, T> {
+    pub fn drain(&mut self) -> Drain<'_, T, A> {
         Drain { iter: self.data.drain(..) }
     }
 
@@ -1438,12 +1493,15 @@ unsafe impl<I> AsVecIntoIter for IntoIter<I> {
 #[must_use = "iterators are lazy and do nothing unless consumed"]
 #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
 #[derive(Clone, Debug)]
-pub struct IntoIterSorted<T> {
-    inner: BinaryHeap<T>,
+pub struct IntoIterSorted<
+    T,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
+> {
+    inner: BinaryHeap<T, A>,
 }
 
 #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
-impl<T: Ord> Iterator for IntoIterSorted<T> {
+impl<T: Ord, A: Allocator> Iterator for IntoIterSorted<T, A> {
     type Item = T;
 
     #[inline]
@@ -1459,13 +1517,13 @@ impl<T: Ord> Iterator for IntoIterSorted<T> {
 }
 
 #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
-impl<T: Ord> ExactSizeIterator for IntoIterSorted<T> {}
+impl<T: Ord, A: Allocator> ExactSizeIterator for IntoIterSorted<T, A> {}
 
 #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
-impl<T: Ord> FusedIterator for IntoIterSorted<T> {}
+impl<T: Ord, A: Allocator> FusedIterator for IntoIterSorted<T, A> {}
 
 #[unstable(feature = "trusted_len", issue = "37572")]
-unsafe impl<T: Ord> TrustedLen for IntoIterSorted<T> {}
+unsafe impl<T: Ord, A: Allocator> TrustedLen for IntoIterSorted<T, A> {}
 
 /// A draining iterator over the elements of a `BinaryHeap`.
 ///
@@ -1475,12 +1533,16 @@ unsafe impl<T: Ord> TrustedLen for IntoIterSorted<T> {}
 /// [`drain`]: BinaryHeap::drain
 #[stable(feature = "drain", since = "1.6.0")]
 #[derive(Debug)]
-pub struct Drain<'a, T: 'a> {
-    iter: vec::Drain<'a, T>,
+pub struct Drain<
+    'a,
+    T: 'a,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + 'a = Global,
+> {
+    iter: vec::Drain<'a, T, A>,
 }
 
 #[stable(feature = "drain", since = "1.6.0")]
-impl<T> Iterator for Drain<'_, T> {
+impl<'a, T, A: Allocator + 'a> Iterator for Drain<'a, T, A> {
     type Item = T;
 
     #[inline]
@@ -1495,22 +1557,22 @@ impl<T> Iterator for Drain<'_, T> {
 }
 
 #[stable(feature = "drain", since = "1.6.0")]
-impl<T> DoubleEndedIterator for Drain<'_, T> {
+impl<'a, T, A: Allocator + 'a> DoubleEndedIterator for Drain<'a, T, A> {
     #[inline]
     fn next_back(&mut self) -> Option<T> {
         self.iter.next_back()
     }
 }
 
 #[stable(feature = "drain", since = "1.6.0")]
-impl<T> ExactSizeIterator for Drain<'_, T> {
+impl<'a, T, A: Allocator + 'a> ExactSizeIterator for Drain<'a, T, A> {
     fn is_empty(&self) -> bool {
         self.iter.is_empty()
     }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
-impl<T> FusedIterator for Drain<'_, T> {}
+impl<'a, T, A: Allocator + 'a> FusedIterator for Drain<'a, T, A> {}
 
 /// A draining iterator over the elements of a `BinaryHeap`.
 ///
@@ -1520,17 +1582,21 @@ impl<T> FusedIterator for Drain<'_, T> {}
 /// [`drain_sorted`]: BinaryHeap::drain_sorted
 #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
 #[derive(Debug)]
-pub struct DrainSorted<'a, T: Ord> {
-    inner: &'a mut BinaryHeap<T>,
+pub struct DrainSorted<
+    'a,
+    T: Ord,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + 'a = Global,
+> {
+    inner: &'a mut BinaryHeap<T, A>,
 }
 
 #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
-impl<'a, T: Ord> Drop for DrainSorted<'a, T> {
+impl<'a, T: Ord, A: Allocator + 'a> Drop for DrainSorted<'a, T, A> {
     /// Removes heap elements in heap order.
     fn drop(&mut self) {
-        struct DropGuard<'r, 'a, T: Ord>(&'r mut DrainSorted<'a, T>);
+        struct DropGuard<'r, 'a, T: Ord, A: Allocator + 'a>(&'r mut DrainSorted<'a, T, A>);
 
-        impl<'r, 'a, T: Ord> Drop for DropGuard<'r, 'a, T> {
+        impl<'r, 'a, T: Ord, A: Allocator + 'a> Drop for DropGuard<'r, 'a, T, A> {
             fn drop(&mut self) {
                 while self.0.inner.pop().is_some() {}
             }
@@ -1545,7 +1611,7 @@ impl<'a, T: Ord> Drop for DrainSorted<'a, T> {
 }
 
 #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
-impl<T: Ord> Iterator for DrainSorted<'_, T> {
+impl<T: Ord, A: Allocator> Iterator for DrainSorted<'_, T, A> {
     type Item = T;
 
     #[inline]
@@ -1561,20 +1627,20 @@ impl<T: Ord> Iterator for DrainSorted<'_, T> {
 }
 
 #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
-impl<T: Ord> ExactSizeIterator for DrainSorted<'_, T> {}
+impl<T: Ord, A: Allocator> ExactSizeIterator for DrainSorted<'_, T, A> {}
 
 #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
-impl<T: Ord> FusedIterator for DrainSorted<'_, T> {}
+impl<T: Ord, A: Allocator> FusedIterator for DrainSorted<'_, T, A> {}
 
 #[unstable(feature = "trusted_len", issue = "37572")]
-unsafe impl<T: Ord> TrustedLen for DrainSorted<'_, T> {}
+unsafe impl<T: Ord, A: Allocator> TrustedLen for DrainSorted<'_, T, A> {}
 
 #[stable(feature = "binary_heap_extras_15", since = "1.5.0")]
-impl<T: Ord> From<Vec<T>> for BinaryHeap<T> {
+impl<T: Ord, A: Allocator> From<Vec<T, A>> for BinaryHeap<T, A> {
     /// Converts a `Vec<T>` into a `BinaryHeap<T>`.
     ///
     /// This conversion happens in-place, and has *O*(*n*) time complexity.
-    fn from(vec: Vec<T>) -> BinaryHeap<T> {
+    fn from(vec: Vec<T, A>) -> BinaryHeap<T, A> {
         let mut heap = BinaryHeap { data: vec };
         heap.rebuild();
         heap
@@ -1598,12 +1664,12 @@ impl<T: Ord, const N: usize> From<[T; N]> for BinaryHeap<T> {
 }
 
 #[stable(feature = "binary_heap_extras_15", since = "1.5.0")]
-impl<T> From<BinaryHeap<T>> for Vec<T> {
+impl<T, A: Allocator> From<BinaryHeap<T, A>> for Vec<T, A> {
     /// Converts a `BinaryHeap<T>` into a `Vec<T>`.
     ///
     /// This conversion requires no data movement or allocation, and has
     /// constant time complexity.
-    fn from(heap: BinaryHeap<T>) -> Vec<T> {
+    fn from(heap: BinaryHeap<T, A>) -> Vec<T, A> {
         heap.data
     }
 }
@@ -1644,7 +1710,7 @@ impl<T> IntoIterator for BinaryHeap<T> {
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> IntoIterator for &'a BinaryHeap<T> {
+impl<'a, T, A: Allocator + 'a> IntoIterator for &'a BinaryHeap<T, A> {
     type Item = &'a T;
     type IntoIter = Iter<'a, T>;
 
@@ -1691,7 +1757,7 @@ impl<T: Ord> SpecExtend<BinaryHeap<T>> for BinaryHeap<T> {
     }
 }
 
-impl<T: Ord> BinaryHeap<T> {
+impl<T: Ord, A: Allocator> BinaryHeap<T, A> {
     fn extend_desugared<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         let iterator = iter.into_iter();
         let (lower, _) = iterator.size_hint();