Skip to content

Commit ce468e6

Browse files
committed
auto merge of #19946 : cgaebel/rust/hashmap-drain-iter, r=gankro
It is useful to move all the elements out of a hashmap without deallocating the underlying buffer. It came up in IRC, and this patch implements it as `drain`. r? @gankro cc: @frankmcsherry
2 parents cc19e33 + d57f259 commit ce468e6

File tree

6 files changed

+470
-18
lines changed

6 files changed

+470
-18
lines changed

src/libcollections/binary_heap.rs

+40-1
Original file line numberDiff line numberDiff line change
@@ -551,9 +551,18 @@ impl<T: Ord> BinaryHeap<T> {
551551
#[unstable = "matches collection reform specification, waiting for dust to settle"]
552552
pub fn is_empty(&self) -> bool { self.len() == 0 }
553553

554+
/// Clears the queue, returning an iterator over the removed elements.
555+
#[inline]
556+
#[unstable = "matches collection reform specification, waiting for dust to settle"]
557+
pub fn drain<'a>(&'a mut self) -> Drain<'a, T> {
558+
Drain {
559+
iter: self.data.drain(),
560+
}
561+
}
562+
554563
/// Drops all items from the queue.
555564
#[unstable = "matches collection reform specification, waiting for dust to settle"]
556-
pub fn clear(&mut self) { self.data.truncate(0) }
565+
pub fn clear(&mut self) { self.drain(); }
557566
}
558567

559568
/// `BinaryHeap` iterator.
@@ -596,6 +605,26 @@ impl<T> DoubleEndedIterator<T> for MoveItems<T> {
596605

597606
impl<T> ExactSizeIterator<T> for MoveItems<T> {}
598607

608+
/// An iterator that drains a `BinaryHeap`.
609+
pub struct Drain<'a, T: 'a> {
610+
iter: vec::Drain<'a, T>,
611+
}
612+
613+
impl<'a, T: 'a> Iterator<T> for Drain<'a, T> {
614+
#[inline]
615+
fn next(&mut self) -> Option<T> { self.iter.next() }
616+
617+
#[inline]
618+
fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
619+
}
620+
621+
impl<'a, T: 'a> DoubleEndedIterator<T> for Drain<'a, T> {
622+
#[inline]
623+
fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
624+
}
625+
626+
impl<'a, T: 'a> ExactSizeIterator<T> for Drain<'a, T> {}
627+
599628
impl<T: Ord> FromIterator<T> for BinaryHeap<T> {
600629
fn from_iter<Iter: Iterator<T>>(iter: Iter) -> BinaryHeap<T> {
601630
let vec: Vec<T> = iter.collect();
@@ -819,4 +848,14 @@ mod tests {
819848
assert_eq!(q.pop().unwrap(), x);
820849
}
821850
}
851+
852+
#[test]
853+
fn test_drain() {
854+
let mut q: BinaryHeap<_> =
855+
[9u, 8, 7, 6, 5, 4, 3, 2, 1].iter().cloned().collect();
856+
857+
assert_eq!(q.drain().take(5).count(), 5);
858+
859+
assert!(q.is_empty());
860+
}
822861
}

src/libcollections/ring_buf.rs

+126-4
Original file line numberDiff line numberDiff line change
@@ -491,6 +491,27 @@ impl<T> RingBuf<T> {
491491
#[unstable = "matches collection reform specification, waiting for dust to settle"]
492492
pub fn is_empty(&self) -> bool { self.len() == 0 }
493493

494+
/// Creates a draining iterator that clears the `RingBuf` and iterates over
495+
/// the removed items from start to end.
496+
///
497+
/// # Examples
498+
///
499+
/// ```
500+
/// use std::collections::RingBuf;
501+
///
502+
/// let mut v = RingBuf::new();
503+
/// v.push_back(1i);
504+
/// assert_eq!(v.drain().next(), Some(1));
505+
/// assert!(v.is_empty());
506+
/// ```
507+
#[inline]
508+
#[unstable = "matches collection reform specification, waiting for dust to settle"]
509+
pub fn drain<'a>(&'a mut self) -> Drain<'a, T> {
510+
Drain {
511+
inner: self,
512+
}
513+
}
514+
494515
/// Clears the buffer, removing all values.
495516
///
496517
/// # Examples
@@ -504,10 +525,9 @@ impl<T> RingBuf<T> {
504525
/// assert!(v.is_empty());
505526
/// ```
506527
#[unstable = "matches collection reform specification, waiting for dust to settle"]
528+
#[inline]
507529
pub fn clear(&mut self) {
508-
while self.pop_front().is_some() {}
509-
self.head = 0;
510-
self.tail = 0;
530+
self.drain();
511531
}
512532

513533
/// Provides a reference to the front element, or `None` if the sequence is
@@ -1230,9 +1250,44 @@ impl<T> DoubleEndedIterator<T> for MoveItems<T> {
12301250
}
12311251
}
12321252

1233-
12341253
impl<T> ExactSizeIterator<T> for MoveItems<T> {}
12351254

1255+
/// A draining RingBuf iterator
1256+
pub struct Drain<'a, T: 'a> {
1257+
inner: &'a mut RingBuf<T>,
1258+
}
1259+
1260+
#[unsafe_destructor]
1261+
impl<'a, T: 'a> Drop for Drain<'a, T> {
1262+
fn drop(&mut self) {
1263+
for _ in *self {}
1264+
self.inner.head = 0;
1265+
self.inner.tail = 0;
1266+
}
1267+
}
1268+
1269+
impl<'a, T: 'a> Iterator<T> for Drain<'a, T> {
1270+
#[inline]
1271+
fn next(&mut self) -> Option<T> {
1272+
self.inner.pop_front()
1273+
}
1274+
1275+
#[inline]
1276+
fn size_hint(&self) -> (uint, Option<uint>) {
1277+
let len = self.inner.len();
1278+
(len, Some(len))
1279+
}
1280+
}
1281+
1282+
impl<'a, T: 'a> DoubleEndedIterator<T> for Drain<'a, T> {
1283+
#[inline]
1284+
fn next_back(&mut self) -> Option<T> {
1285+
self.inner.pop_back()
1286+
}
1287+
}
1288+
1289+
impl<'a, T: 'a> ExactSizeIterator<T> for Drain<'a, T> {}
1290+
12361291
impl<A: PartialEq> PartialEq for RingBuf<A> {
12371292
fn eq(&self, other: &RingBuf<A>) -> bool {
12381293
self.len() == other.len() &&
@@ -1841,6 +1896,73 @@ mod tests {
18411896
}
18421897
}
18431898

1899+
#[test]
1900+
fn test_drain() {
1901+
1902+
// Empty iter
1903+
{
1904+
let mut d: RingBuf<int> = RingBuf::new();
1905+
1906+
{
1907+
let mut iter = d.drain();
1908+
1909+
assert_eq!(iter.size_hint(), (0, Some(0)));
1910+
assert_eq!(iter.next(), None);
1911+
assert_eq!(iter.size_hint(), (0, Some(0)));
1912+
}
1913+
1914+
assert!(d.is_empty());
1915+
}
1916+
1917+
// simple iter
1918+
{
1919+
let mut d = RingBuf::new();
1920+
for i in range(0i, 5) {
1921+
d.push_back(i);
1922+
}
1923+
1924+
assert_eq!(d.drain().collect::<Vec<int>>(), [0, 1, 2, 3, 4]);
1925+
assert!(d.is_empty());
1926+
}
1927+
1928+
// wrapped iter
1929+
{
1930+
let mut d = RingBuf::new();
1931+
for i in range(0i, 5) {
1932+
d.push_back(i);
1933+
}
1934+
for i in range(6, 9) {
1935+
d.push_front(i);
1936+
}
1937+
1938+
assert_eq!(d.drain().collect::<Vec<int>>(), [8,7,6,0,1,2,3,4]);
1939+
assert!(d.is_empty());
1940+
}
1941+
1942+
// partially used
1943+
{
1944+
let mut d = RingBuf::new();
1945+
for i in range(0i, 5) {
1946+
d.push_back(i);
1947+
}
1948+
for i in range(6, 9) {
1949+
d.push_front(i);
1950+
}
1951+
1952+
{
1953+
let mut it = d.drain();
1954+
assert_eq!(it.size_hint(), (8, Some(8)));
1955+
assert_eq!(it.next(), Some(8));
1956+
assert_eq!(it.size_hint(), (7, Some(7)));
1957+
assert_eq!(it.next_back(), Some(4));
1958+
assert_eq!(it.size_hint(), (6, Some(6)));
1959+
assert_eq!(it.next(), Some(7));
1960+
assert_eq!(it.size_hint(), (5, Some(5)));
1961+
}
1962+
assert!(d.is_empty());
1963+
}
1964+
}
1965+
18441966
#[test]
18451967
fn test_from_iter() {
18461968
use core::iter;

0 commit comments

Comments
 (0)