From 2a34a7b83914f3ff7d2d52827b4c92bffefb5f92 Mon Sep 17 00:00:00 2001
From: Nathan Moos <moosingin3space@gmail.com>
Date: Mon, 23 May 2016 16:59:11 -0700
Subject: [PATCH] implemented peek_mut and unit tests

---
 src/libcollections/binary_heap.rs     | 68 +++++++++++++++++++++++++++
 src/libcollectionstest/binary_heap.rs | 18 +++++++
 src/libcollectionstest/lib.rs         |  1 +
 3 files changed, 87 insertions(+)

diff --git a/src/libcollections/binary_heap.rs b/src/libcollections/binary_heap.rs
index 43c6e6e81209d..140801737bc4c 100644
--- a/src/libcollections/binary_heap.rs
+++ b/src/libcollections/binary_heap.rs
@@ -151,6 +151,7 @@
 #![allow(missing_docs)]
 #![stable(feature = "rust1", since = "1.0.0")]
 
+use core::ops::{Drop, Deref, DerefMut};
 use core::iter::FromIterator;
 use core::mem::swap;
 use core::mem::size_of;
@@ -218,6 +219,37 @@ pub struct BinaryHeap<T> {
     data: Vec<T>,
 }
 
+/// A container object that represents the result of the [`peek_mut()`] method
+/// on `BinaryHeap`. See its documentation for details.
+///
+/// [`peek_mut()`]: struct.BinaryHeap.html#method.peek_mut
+#[unstable(feature = "binary_heap_peek_mut", issue = "34392")]
+pub struct PeekMut<'a, T: 'a + Ord> {
+    heap: &'a mut BinaryHeap<T>
+}
+
+#[unstable(feature = "binary_heap_peek_mut", issue = "34392")]
+impl<'a, T: Ord> Drop for PeekMut<'a, T> {
+    fn drop(&mut self) {
+        self.heap.sift_down(0);
+    }
+}
+
+#[unstable(feature = "binary_heap_peek_mut", issue = "34392")]
+impl<'a, T: Ord> Deref for PeekMut<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        &self.heap.data[0]
+    }
+}
+
+#[unstable(feature = "binary_heap_peek_mut", issue = "34392")]
+impl<'a, T: Ord> DerefMut for PeekMut<'a, T> {
+    fn deref_mut(&mut self) -> &mut T {
+        &mut self.heap.data[0]
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T: Clone> Clone for BinaryHeap<T> {
     fn clone(&self) -> Self {
@@ -323,6 +355,42 @@ impl<T: Ord> BinaryHeap<T> {
         self.data.get(0)
     }
 
+    /// Returns a mutable reference to the greatest item in the binary heap, or
+    /// `None` if it is empty.
+    ///
+    /// Note: If the `PeekMut` value is leaked, the heap may be in an
+    /// inconsistent state.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(binary_heap_peek_mut)]
+    /// use std::collections::BinaryHeap;
+    /// let mut heap = BinaryHeap::new();
+    /// assert!(heap.peek_mut().is_none());
+    ///
+    /// heap.push(1);
+    /// heap.push(5);
+    /// heap.push(2);
+    /// {
+    ///     let mut val = heap.peek_mut().unwrap();
+    ///     *val = 0;
+    /// }
+    /// assert_eq!(heap.peek(), Some(&2));
+    /// ```
+    #[unstable(feature = "binary_heap_peek_mut", issue = "34392")]
+    pub fn peek_mut(&mut self) -> Option<PeekMut<T>> {
+        if self.is_empty() {
+            None
+        } else {
+            Some(PeekMut {
+                heap: self
+            })
+        }
+    }
+
     /// Returns the number of elements the binary heap can hold without reallocating.
     ///
     /// # Examples
diff --git a/src/libcollectionstest/binary_heap.rs b/src/libcollectionstest/binary_heap.rs
index 58194fe75f79a..be933abe41fe2 100644
--- a/src/libcollectionstest/binary_heap.rs
+++ b/src/libcollectionstest/binary_heap.rs
@@ -81,6 +81,18 @@ fn test_peek_and_pop() {
     }
 }
 
+#[test]
+fn test_peek_mut() {
+    let data = vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1];
+    let mut heap = BinaryHeap::from(data);
+    assert_eq!(heap.peek(), Some(&10));
+    {
+        let mut top = heap.peek_mut().unwrap();
+        *top -= 2;
+    }
+    assert_eq!(heap.peek(), Some(&9));
+}
+
 #[test]
 fn test_push() {
     let mut heap = BinaryHeap::from(vec![2, 4, 9]);
@@ -192,6 +204,12 @@ fn test_empty_peek() {
     assert!(empty.peek().is_none());
 }
 
+#[test]
+fn test_empty_peek_mut() {
+    let mut empty = BinaryHeap::<i32>::new();
+    assert!(empty.peek_mut().is_none());
+}
+
 #[test]
 fn test_empty_replace() {
     let mut heap = BinaryHeap::new();
diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs
index 400d614094862..6161bad7468c6 100644
--- a/src/libcollectionstest/lib.rs
+++ b/src/libcollectionstest/lib.rs
@@ -12,6 +12,7 @@
 
 #![feature(binary_heap_extras)]
 #![feature(binary_heap_append)]
+#![feature(binary_heap_peek_mut)]
 #![feature(box_syntax)]
 #![feature(btree_append)]
 #![feature(btree_split_off)]