diff --git a/src/libcollections/binary_heap.rs b/src/libcollections/binary_heap.rs index c5d5ad27d23..c53b34d589b 100644 --- a/src/libcollections/binary_heap.rs +++ b/src/libcollections/binary_heap.rs @@ -153,8 +153,7 @@ use core::ops::{Deref, DerefMut}; use core::iter::{FromIterator, FusedIterator}; -use core::mem::swap; -use core::mem::size_of; +use core::mem::{swap, size_of}; use core::ptr; use core::fmt; @@ -226,12 +225,15 @@ pub struct BinaryHeap { #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")] pub struct PeekMut<'a, T: 'a + Ord> { heap: &'a mut BinaryHeap, + sift: bool, } #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")] impl<'a, T: Ord> Drop for PeekMut<'a, T> { fn drop(&mut self) { - self.heap.sift_down(0); + if self.sift { + self.heap.sift_down(0); + } } } @@ -250,6 +252,16 @@ impl<'a, T: Ord> DerefMut for PeekMut<'a, T> { } } +impl<'a, T: Ord> PeekMut<'a, T> { + /// Removes the peeked value from the heap and returns it. + #[unstable(feature = "binary_heap_peek_mut_pop", issue = "0")] + pub fn pop(mut this: PeekMut<'a, T>) -> T { + let value = this.heap.pop().unwrap(); + this.sift = false; + value + } +} + #[stable(feature = "rust1", since = "1.0.0")] impl Clone for BinaryHeap { fn clone(&self) -> Self { @@ -385,7 +397,10 @@ impl BinaryHeap { if self.is_empty() { None } else { - Some(PeekMut { heap: self }) + Some(PeekMut { + heap: self, + sift: true, + }) } } diff --git a/src/libcollectionstest/binary_heap.rs b/src/libcollectionstest/binary_heap.rs index 9cd63d87931..1df341d1fc2 100644 --- a/src/libcollectionstest/binary_heap.rs +++ b/src/libcollectionstest/binary_heap.rs @@ -9,7 +9,7 @@ // except according to those terms. use std::collections::BinaryHeap; -use std::collections::binary_heap::Drain; +use std::collections::binary_heap::{Drain, PeekMut}; #[test] fn test_iterator() { @@ -94,6 +94,19 @@ fn test_peek_mut() { assert_eq!(heap.peek(), Some(&9)); } +#[test] +fn test_peek_mut_pop() { + 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!(PeekMut::pop(top), 8); + } + assert_eq!(heap.peek(), Some(&9)); +} + #[test] fn test_push() { let mut heap = BinaryHeap::from(vec![2, 4, 9]); diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs index d4fb5ea03ad..ee555ee351e 100644 --- a/src/libcollectionstest/lib.rs +++ b/src/libcollectionstest/lib.rs @@ -11,6 +11,7 @@ #![deny(warnings)] #![feature(binary_heap_extras)] +#![feature(binary_heap_peek_mut_pop)] #![feature(box_syntax)] #![feature(btree_range)] #![feature(collections)]