Skip to content

Commit

Permalink
Add PeekMut::pop
Browse files Browse the repository at this point in the history
A fairly common workflow is to put a bunch of stuff into a binary heap
and then mutate the top value until its empty. This both makes that a
bit more convenient (no need to save a boolean off and pop after to
avoid borrowck issues), and a bit more efficient since you only shift
once.
  • Loading branch information
sfackler committed Jan 2, 2017
1 parent 6c9bb42 commit f2cb47a
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 5 deletions.
23 changes: 19 additions & 4 deletions src/libcollections/binary_heap.rs
Expand Up @@ -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;

Expand Down Expand Up @@ -226,12 +225,15 @@ pub struct BinaryHeap<T> {
#[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
pub struct PeekMut<'a, T: 'a + Ord> {
heap: &'a mut BinaryHeap<T>,
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);
}
}
}

Expand All @@ -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<T: Clone> Clone for BinaryHeap<T> {
fn clone(&self) -> Self {
Expand Down Expand Up @@ -385,7 +397,10 @@ impl<T: Ord> BinaryHeap<T> {
if self.is_empty() {
None
} else {
Some(PeekMut { heap: self })
Some(PeekMut {
heap: self,
sift: true,
})
}
}

Expand Down
15 changes: 14 additions & 1 deletion src/libcollectionstest/binary_heap.rs
Expand Up @@ -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() {
Expand Down Expand Up @@ -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]);
Expand Down
1 change: 1 addition & 0 deletions src/libcollectionstest/lib.rs
Expand Up @@ -11,6 +11,7 @@
#![deny(warnings)]

#![feature(binary_heap_extras)]
#![feature(binary_heap_peek_mut_pop)]
#![feature(box_syntax)]
#![feature(btree_range)]
#![feature(collections)]
Expand Down

0 comments on commit f2cb47a

Please sign in to comment.