Implement drain iterator (#56)

src/lib.rs

@@ -104,6 +104,7 @@
 
 use std::iter::IntoIterator;
 use std::ops;
+use std::vec;
 use std::{fmt, mem};
 
 /// Pre-allocated storage for a uniform data type
@@ -170,6 +171,9 @@ pub struct IterMut<'a, T: 'a> {
     curr: usize,
 }
 
+/// A draining iterator for `Slab`
+pub struct Drain<'a, T: 'a>(vec::Drain<'a, Entry<T>>);
+
 #[derive(Clone)]
 enum Entry<T> {
     Vacant(usize),
@@ -741,6 +745,39 @@ impl<T> Slab<T> {
             }
         }
     }
+
+    /// Return a draining iterator that removes all elements from the slab and
+    /// yields the removed items.
+    ///
+    /// Note: Elements are removed even if the iterator is only partially
+    /// consumed or not consumed at all.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use slab::*;
+    /// let mut slab = Slab::new();
+    ///
+    /// let _ = slab.insert(0);
+    /// let _ = slab.insert(1);
+    /// let _ = slab.insert(2);
+    ///
+    /// {
+    ///     let mut drain = slab.drain();
+    ///
+    ///     assert_eq!(Some(0), drain.next());
+    ///     assert_eq!(Some(1), drain.next());
+    ///     assert_eq!(Some(2), drain.next());
+    ///     assert_eq!(None, drain.next());
+    /// }
+    ///
+    /// assert!(slab.is_empty());
+    /// ```
+    pub fn drain(&mut self) -> Drain<T> {
+        self.len = 0;
+        self.next = 0;
+        Drain(self.entries.drain(..))
+    }
 }
 
 impl<T> ops::Index<usize> for Slab<T> {
@@ -819,6 +856,12 @@ where
     }
 }
 
+impl<'a, T: 'a> fmt::Debug for Drain<'a, T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("Drain").finish()
+    }
+}
+
 // ===== VacantEntry =====
 
 impl<'a, T> VacantEntry<'a, T> {
@@ -916,3 +959,19 @@ impl<'a, T> Iterator for IterMut<'a, T> {
         None
     }
 }
+
+// ===== Drain =====
+
+impl<'a, T> Iterator for Drain<'a, T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        while let Some(entry) = self.0.next() {
+            if let Entry::Occupied(v) = entry {
+                return Some(v);
+            }
+        }
+
+        None
+    }
+}

tests/slab.rs

@@ -264,3 +264,38 @@ fn clear() {
     let vals: Vec<_> = slab.iter().map(|(_, r)| *r).collect();
     assert!(vals.is_empty());
 }
+
+#[test]
+fn fully_consumed_drain() {
+    let mut slab = Slab::new();
+
+    for i in 0..3 {
+        slab.insert(i);
+    }
+
+    {
+        let mut drain = slab.drain();
+        assert_eq!(Some(0), drain.next());
+        assert_eq!(Some(1), drain.next());
+        assert_eq!(Some(2), drain.next());
+        assert_eq!(None, drain.next());
+    }
+
+    assert!(slab.is_empty());
+}
+
+#[test]
+fn partially_consumed_drain() {
+    let mut slab = Slab::new();
+
+    for i in 0..3 {
+        slab.insert(i);
+    }
+
+    {
+        let mut drain = slab.drain();
+        assert_eq!(Some(0), drain.next());
+    }
+
+    assert!(slab.is_empty())
+}