Rollup merge of #46262 - udoprog:linked-list-remove-if, r=dtolnay

Introduce LinkedList::drain_filter

This introduces `LinkedList::remove_if`.

This operation embodies one of the use-cases where `LinkedList` would typically be preferred over `Vec`: random removal and retrieval.

There are a number of considerations with this:

Should there be two `remove_if` methods? One where elements are only removed, one which returns a collection of removed elements.

Should this be implemented using a draining iterator pattern that covers both cases? I suspect that would incur a bit of overhead (moving the element into the iterator, then into a new collection). But I'm not sure. Maybe that's an acceptable compromise to maximize flexibility.

I don't feel I've had enough exposure to unsafe programming in rust to be certain the implementation is correct. This relies quite heavily on moving around copies of Shared pointers to make the code reasonable. Please help me out :).

2 files changed, 327 insertions, 0 deletions

diff --git a/src/liballoc/linked_list.rs b/src/liballoc/linked_list.rs
index 99ad424cc20..0fe3c972422 100644
--- a/src/liballoc/linked_list.rs
+++ b/src/liballoc/linked_list.rs
@@ -220,6 +220,28 @@ impl<T> LinkedList<T> {
             node
         })
     }
+
+    /// Unlinks the specified node from the current list.
+    ///
+    /// Warning: this will not check that the provided node belongs to the current list.
+    #[inline]
+    unsafe fn unlink_node(&mut self, mut node: Shared<Node<T>>) {
+        let node = node.as_mut();
+
+        match node.prev {
+            Some(mut prev) => prev.as_mut().next = node.next.clone(),
+            // this node is the head node
+            None => self.head = node.next.clone(),
+        };
+
+        match node.next {
+            Some(mut next) => next.as_mut().prev = node.prev.clone(),
+            // this node is the tail node
+            None => self.tail = node.prev.clone(),
+        };
+
+        self.len -= 1;
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -722,6 +744,49 @@ impl<T> LinkedList<T> {
         second_part
     }
 
+    /// Creates an iterator which uses a closure to determine if an element should be removed.
+    ///
+    /// 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
+    ///
+    /// Splitting a list into evens and odds, reusing the original list:
+    ///
+    /// ```
+    /// #![feature(drain_filter)]
+    /// use std::collections::LinkedList;
+    ///
+    /// 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 = "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
+    {
+        // avoid borrow issues.
+        let it = self.head;
+        let old_len = self.len;
+
+        DrainFilter {
+            list: self,
+            it: it,
+            pred: filter,
+            idx: 0,
+            old_len: old_len,
+        }
+    }
+
     /// Returns a place for insertion at the front of the list.
     ///
     /// Using this method with placement syntax is equivalent to
@@ -967,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;
@@ -1509,4 +1624,28 @@ mod tests {
         }
         assert_eq!(i, v.len());
     }
+
+    #[test]
+    fn drain_filter_test() {
+        let mut m: LinkedList<u32> = LinkedList::new();
+        m.extend(&[1, 2, 3, 4, 5, 6]);
+        let deleted = m.drain_filter(|v| *v < 4).collect::<Vec<_>>();
+
+        check_links(&m);
+
+        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]);
+    }
+}