Skip to content

Commit d57f259

Browse files
author
Clark Gaebel
committed
[collections] Adds drain: a way to sneak out the elements while clearing.
It is useful to move all the elements out of some collections without deallocating the underlying buffer. It came up in IRC, and this patch implements it as `drain`. This has been discussed as part of RFC 509. r? @gankro cc: @frankmcsherry
1 parent f9a4849 commit d57f259

File tree

6 files changed

+470
-18
lines changed

6 files changed

+470
-18
lines changed

src/libcollections/binary_heap.rs

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

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

561570
/// `BinaryHeap` iterator.
@@ -598,6 +607,26 @@ impl<T> DoubleEndedIterator<T> for MoveItems<T> {
598607

599608
impl<T> ExactSizeIterator<T> for MoveItems<T> {}
600609

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

src/libcollections/ring_buf.rs

Lines changed: 126 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -443,6 +443,27 @@ impl<T> RingBuf<T> {
443443
#[unstable = "matches collection reform specification, waiting for dust to settle"]
444444
pub fn is_empty(&self) -> bool { self.len() == 0 }
445445

446+
/// Creates a draining iterator that clears the `RingBuf` and iterates over
447+
/// the removed items from start to end.
448+
///
449+
/// # Examples
450+
///
451+
/// ```
452+
/// use std::collections::RingBuf;
453+
///
454+
/// let mut v = RingBuf::new();
455+
/// v.push_back(1i);
456+
/// assert_eq!(v.drain().next(), Some(1));
457+
/// assert!(v.is_empty());
458+
/// ```
459+
#[inline]
460+
#[unstable = "matches collection reform specification, waiting for dust to settle"]
461+
pub fn drain<'a>(&'a mut self) -> Drain<'a, T> {
462+
Drain {
463+
inner: self,
464+
}
465+
}
466+
446467
/// Clears the buffer, removing all values.
447468
///
448469
/// # Examples
@@ -456,10 +477,9 @@ impl<T> RingBuf<T> {
456477
/// assert!(v.is_empty());
457478
/// ```
458479
#[unstable = "matches collection reform specification, waiting for dust to settle"]
480+
#[inline]
459481
pub fn clear(&mut self) {
460-
while self.pop_front().is_some() {}
461-
self.head = 0;
462-
self.tail = 0;
482+
self.drain();
463483
}
464484

465485
/// Provides a reference to the front element, or `None` if the sequence is
@@ -1177,9 +1197,44 @@ impl<T> DoubleEndedIterator<T> for MoveItems<T> {
11771197
}
11781198
}
11791199

1180-
11811200
impl<T> ExactSizeIterator<T> for MoveItems<T> {}
11821201

1202+
/// A draining RingBuf iterator
1203+
pub struct Drain<'a, T: 'a> {
1204+
inner: &'a mut RingBuf<T>,
1205+
}
1206+
1207+
#[unsafe_destructor]
1208+
impl<'a, T: 'a> Drop for Drain<'a, T> {
1209+
fn drop(&mut self) {
1210+
for _ in *self {}
1211+
self.inner.head = 0;
1212+
self.inner.tail = 0;
1213+
}
1214+
}
1215+
1216+
impl<'a, T: 'a> Iterator<T> for Drain<'a, T> {
1217+
#[inline]
1218+
fn next(&mut self) -> Option<T> {
1219+
self.inner.pop_front()
1220+
}
1221+
1222+
#[inline]
1223+
fn size_hint(&self) -> (uint, Option<uint>) {
1224+
let len = self.inner.len();
1225+
(len, Some(len))
1226+
}
1227+
}
1228+
1229+
impl<'a, T: 'a> DoubleEndedIterator<T> for Drain<'a, T> {
1230+
#[inline]
1231+
fn next_back(&mut self) -> Option<T> {
1232+
self.inner.pop_back()
1233+
}
1234+
}
1235+
1236+
impl<'a, T: 'a> ExactSizeIterator<T> for Drain<'a, T> {}
1237+
11831238
impl<A: PartialEq> PartialEq for RingBuf<A> {
11841239
fn eq(&self, other: &RingBuf<A>) -> bool {
11851240
self.len() == other.len() &&
@@ -1789,6 +1844,73 @@ mod tests {
17891844
}
17901845
}
17911846

1847+
#[test]
1848+
fn test_drain() {
1849+
1850+
// Empty iter
1851+
{
1852+
let mut d: RingBuf<int> = RingBuf::new();
1853+
1854+
{
1855+
let mut iter = d.drain();
1856+
1857+
assert_eq!(iter.size_hint(), (0, Some(0)));
1858+
assert_eq!(iter.next(), None);
1859+
assert_eq!(iter.size_hint(), (0, Some(0)));
1860+
}
1861+
1862+
assert!(d.is_empty());
1863+
}
1864+
1865+
// simple iter
1866+
{
1867+
let mut d = RingBuf::new();
1868+
for i in range(0i, 5) {
1869+
d.push_back(i);
1870+
}
1871+
1872+
assert_eq!(d.drain().collect::<Vec<int>>(), [0, 1, 2, 3, 4]);
1873+
assert!(d.is_empty());
1874+
}
1875+
1876+
// wrapped iter
1877+
{
1878+
let mut d = RingBuf::new();
1879+
for i in range(0i, 5) {
1880+
d.push_back(i);
1881+
}
1882+
for i in range(6, 9) {
1883+
d.push_front(i);
1884+
}
1885+
1886+
assert_eq!(d.drain().collect::<Vec<int>>(), [8,7,6,0,1,2,3,4]);
1887+
assert!(d.is_empty());
1888+
}
1889+
1890+
// partially used
1891+
{
1892+
let mut d = RingBuf::new();
1893+
for i in range(0i, 5) {
1894+
d.push_back(i);
1895+
}
1896+
for i in range(6, 9) {
1897+
d.push_front(i);
1898+
}
1899+
1900+
{
1901+
let mut it = d.drain();
1902+
assert_eq!(it.size_hint(), (8, Some(8)));
1903+
assert_eq!(it.next(), Some(8));
1904+
assert_eq!(it.size_hint(), (7, Some(7)));
1905+
assert_eq!(it.next_back(), Some(4));
1906+
assert_eq!(it.size_hint(), (6, Some(6)));
1907+
assert_eq!(it.next(), Some(7));
1908+
assert_eq!(it.size_hint(), (5, Some(5)));
1909+
}
1910+
assert!(d.is_empty());
1911+
}
1912+
}
1913+
17921914
#[test]
17931915
fn test_from_iter() {
17941916
use std::iter;

0 commit comments

Comments
 (0)