Implement LinkedList::drain_filter

Relates to rust-lang/rfcs#2140 - drain_filter for all collections

`drain_filter` is implemented instead of `LinkedList::remove_if` based
on review feedback.

2 files changed, 283 insertions, 35 deletions

diff --git a/src/liballoc/linked_list.rs b/src/liballoc/linked_list.rs
index 689a2671994..7bd1e1f0504 100644
--- a/src/liballoc/linked_list.rs
+++ b/src/liballoc/linked_list.rs
@@ -744,48 +744,47 @@ impl<T> LinkedList<T> {
         second_part
     }
 
-    /// Removes any element matching the given predicate. Returns the elements which were removed
-    /// in a new list.
+    /// Creates an iterator which uses a closure to determine if an element should be removed.
     ///
-    /// This operation should compute in O(n) time.
+    /// If the closure returns true, then the element is removed and yielded.
+    /// If the closure returns false, it will try again, and call the closure on the next element,
+    /// seeing if it passes the test.
+    ///
+    /// Note that `drain_filter` lets you mutate every element in the filter closure, regardless of
+    /// whether you choose to keep or remove it.
     ///
     /// # Examples
     ///
-    /// ```
-    /// #![feature(linked_list_remove_if)]
+    /// Splitting a list into evens and odds, reusing the original list:
     ///
+    /// ```
+    /// #![feature(drain_filter)]
     /// use std::collections::LinkedList;
     ///
-    /// let mut d = LinkedList::new();
-    /// d.push_back(1);
-    /// d.push_back(2);
-    /// d.push_back(3);
-    /// assert_eq!(d.remove_if(|v| *v < 3).len(), 2);
-    /// assert_eq!(d.len(), 1);
+    /// let mut numbers: LinkedList<u32> = LinkedList::new();
+    /// numbers.extend(&[1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
+    ///
+    /// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<LinkedList<_>>();
+    /// let odds = numbers;
+    ///
+    /// assert_eq!(evens.into_iter().collect::<Vec<_>>(), vec![2, 4, 6, 8, 14]);
+    /// assert_eq!(odds.into_iter().collect::<Vec<_>>(), vec![1, 3, 5, 9, 11, 13, 15]);
     /// ```
-    #[unstable(feature = "linked_list_remove_if",
-               reason = "experimental method",
-               issue = "0")]
-    pub fn remove_if<P>(&mut self, predicate: P) -> LinkedList<T>
-        where P: Fn(&T) -> bool
+    #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+    pub fn drain_filter<F>(&mut self, filter: F) -> DrainFilter<T, F>
+        where F: FnMut(&mut T) -> bool
     {
-        let mut deleted = LinkedList::new();
-
-        let mut it = self.head;
-
-        while let Some(node) = it {
-            unsafe {
-                it = node.as_ref().next;
+        // avoid borrow issues.
+        let it = self.head;
+        let old_len = self.len;
 
-                if predicate(&node.as_ref().element) {
-                    self.unlink_node(node);
-                    // move the unlinked node into the deleted list.
-                    deleted.push_back_node(Box::from_raw(node.as_ptr()));
-                }
-            }
+        DrainFilter {
+            list: self,
+            it: it,
+            pred: filter,
+            idx: 0,
+            old_len: old_len,
         }
-
-        deleted
     }
 
     /// Returns a place for insertion at the front of the list.
@@ -1033,6 +1032,56 @@ impl<'a, T> IterMut<'a, T> {
     }
 }
 
+/// An iterator produced by calling `drain_filter` on LinkedList.
+#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+pub struct DrainFilter<'a, T: 'a, F: 'a>
+    where F: FnMut(&mut T) -> bool,
+{
+    list: &'a mut LinkedList<T>,
+    it: Option<Shared<Node<T>>>,
+    pred: F,
+    idx: usize,
+    old_len: usize,
+}
+
+#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+impl<'a, T, F> Iterator for DrainFilter<'a, T, F>
+    where F: FnMut(&mut T) -> bool,
+{
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        while let Some(mut node) = self.it {
+            unsafe {
+                self.it = node.as_ref().next;
+                self.idx += 1;
+
+                if (self.pred)(&mut node.as_mut().element) {
+                    self.list.unlink_node(node);
+                    return Some(Box::from_raw(node.as_ptr()).element);
+                }
+            }
+        }
+
+        None
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (0, Some(self.old_len - self.idx))
+    }
+}
+
+#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+impl<'a, T: 'a + fmt::Debug, F> fmt::Debug for DrainFilter<'a, T, F>
+    where F: FnMut(&mut T) -> bool
+{
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_tuple("DrainFilter")
+         .field(&self.list)
+         .finish()
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T> Iterator for IntoIter<T> {
     type Item = T;
@@ -1570,15 +1619,26 @@ mod tests {
     }
 
     #[test]
-    fn remove_if_test() {
+    fn drain_filter_test() {
         let mut m: LinkedList<u32> = LinkedList::new();
         m.extend(&[1, 2, 3, 4, 5, 6]);
-        let deleted = m.remove_if(|v| *v < 4);
+        let deleted = m.drain_filter(|v| *v < 4).collect::<Vec<_>>();
 
         check_links(&m);
-        check_links(&deleted);
 
-        assert_eq!(deleted.into_iter().collect::<Vec<_>>(), &[1, 2, 3]);
+        assert_eq!(deleted, &[1, 2, 3]);
         assert_eq!(m.into_iter().collect::<Vec<_>>(), &[4, 5, 6]);
     }
+
+    #[test]
+    fn drain_to_empty_test() {
+        let mut m: LinkedList<u32> = LinkedList::new();
+        m.extend(&[1, 2, 3, 4, 5, 6]);
+        let deleted = m.drain_filter(|_| true).collect::<Vec<_>>();
+
+        check_links(&m);
+
+        assert_eq!(deleted, &[1, 2, 3, 4, 5, 6]);
+        assert_eq!(m.into_iter().collect::<Vec<_>>(), &[]);
+    }
 }
diff --git a/src/liballoc/tests/linked_list.rs b/src/liballoc/tests/linked_list.rs
index a59724a017b..4e3e855105e 100644
--- a/src/liballoc/tests/linked_list.rs
+++ b/src/liballoc/tests/linked_list.rs
@@ -366,3 +366,191 @@ fn test_contains() {
 
     assert!(!l.contains(&3));
 }
+
+#[test]
+fn drain_filter_empty() {
+    let mut list: LinkedList<i32> = LinkedList::new();
+
+    {
+        let mut iter = list.drain_filter(|_| true);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(list.len(), 0);
+    assert_eq!(list.into_iter().collect::<Vec<_>>(), vec![]);
+}
+
+#[test]
+fn drain_filter_zst() {
+    let mut list: LinkedList<_> = vec![(), (), (), (), ()].into_iter().collect();
+    let initial_len = list.len();
+    let mut count = 0;
+
+    {
+        let mut iter = list.drain_filter(|_| true);
+        assert_eq!(iter.size_hint(), (0, Some(initial_len)));
+        while let Some(_) = iter.next() {
+            count += 1;
+            assert_eq!(iter.size_hint(), (0, Some(initial_len - count)));
+        }
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(count, initial_len);
+    assert_eq!(list.len(), 0);
+    assert_eq!(list.into_iter().collect::<Vec<_>>(), vec![]);
+}
+
+#[test]
+fn drain_filter_false() {
+    let mut list: LinkedList<_> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_iter().collect();
+
+    let initial_len = list.len();
+    let mut count = 0;
+
+    {
+        let mut iter = list.drain_filter(|_| false);
+        assert_eq!(iter.size_hint(), (0, Some(initial_len)));
+        for _ in iter.by_ref() {
+            count += 1;
+        }
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(count, 0);
+    assert_eq!(list.len(), initial_len);
+    assert_eq!(list.into_iter().collect::<Vec<_>>(), vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+}
+
+#[test]
+fn drain_filter_true() {
+    let mut list: LinkedList<_> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_iter().collect();
+
+    let initial_len = list.len();
+    let mut count = 0;
+
+    {
+        let mut iter = list.drain_filter(|_| true);
+        assert_eq!(iter.size_hint(), (0, Some(initial_len)));
+        while let Some(_) = iter.next() {
+            count += 1;
+            assert_eq!(iter.size_hint(), (0, Some(initial_len - count)));
+        }
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(count, initial_len);
+    assert_eq!(list.len(), 0);
+    assert_eq!(list.into_iter().collect::<Vec<_>>(), vec![]);
+}
+
+#[test]
+fn drain_filter_complex() {
+
+    {   //                [+xxx++++++xxxxx++++x+x++]
+        let mut list = vec![
+            1,
+            2, 4, 6,
+            7, 9, 11, 13, 15, 17,
+            18, 20, 22, 24, 26,
+            27, 29, 31, 33,
+            34,
+            35,
+            36,
+            37, 39
+        ].into_iter().collect::<LinkedList<_>>();
+
+        let removed = list.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);
+
+        assert_eq!(list.len(), 14);
+        assert_eq!(
+            list.into_iter().collect::<Vec<_>>(),
+            vec![1, 7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35, 37, 39]
+        );
+    }
+
+    {   // [xxx++++++xxxxx++++x+x++]
+        let mut list = vec![
+            2, 4, 6,
+            7, 9, 11, 13, 15, 17,
+            18, 20, 22, 24, 26,
+            27, 29, 31, 33,
+            34,
+            35,
+            36,
+            37, 39
+        ].into_iter().collect::<LinkedList<_>>();
+
+        let removed = list.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);
+
+        assert_eq!(list.len(), 13);
+        assert_eq!(
+            list.into_iter().collect::<Vec<_>>(),
+            vec![7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35, 37, 39]
+        );
+    }
+
+    {   // [xxx++++++xxxxx++++x+x]
+        let mut list = vec![
+            2, 4, 6,
+            7, 9, 11, 13, 15, 17,
+            18, 20, 22, 24, 26,
+            27, 29, 31, 33,
+            34,
+            35,
+            36
+        ].into_iter().collect::<LinkedList<_>>();
+
+        let removed = list.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);
+
+        assert_eq!(list.len(), 11);
+        assert_eq!(
+            list.into_iter().collect::<Vec<_>>(),
+            vec![7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35]
+        );
+    }
+
+    {   // [xxxxxxxxxx+++++++++++]
+        let mut list = vec![
+            2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
+            1, 3, 5, 7, 9, 11, 13, 15, 17, 19
+        ].into_iter().collect::<LinkedList<_>>();
+
+        let removed = list.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]);
+
+        assert_eq!(list.len(), 10);
+        assert_eq!(list.into_iter().collect::<Vec<_>>(), vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
+    }
+
+    {   // [+++++++++++xxxxxxxxxx]
+        let mut list = vec![
+            1, 3, 5, 7, 9, 11, 13, 15, 17, 19,
+            2, 4, 6, 8, 10, 12, 14, 16, 18, 20
+        ].into_iter().collect::<LinkedList<_>>();
+
+        let removed = list.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]);
+
+        assert_eq!(list.len(), 10);
+        assert_eq!(list.into_iter().collect::<Vec<_>>(), vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
+    }
+}