From c5589aecdf748a682eb8080180049c6559ea85f7 Mon Sep 17 00:00:00 2001
From: bluss <bluss@users.noreply.github.com>
Date: Wed, 17 Feb 2016 23:08:46 +0100
Subject: [PATCH] Add .retain() method

Based on "retain_mut" as discussed for vec: we can provide &mut to the
element to remove effortlessly.
---
 src/lib.rs     | 34 ++++++++++++++++++++++++++++++++++
 tests/tests.rs | 17 +++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 9fd9886..11c710a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -277,6 +277,40 @@ impl<A: Array> ArrayVec<A> {
         while let Some(_) = self.pop() { }
     }
 
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// In other words, remove all elements `e` such that `f(&mut e)` returns false.
+    /// This method operates in place and preserves the order of the retained
+    /// elements.
+    ///
+    /// ```
+    /// use arrayvec::ArrayVec;
+    ///
+    /// let mut array = ArrayVec::from([1, 2, 3, 4]);
+    /// array.retain(|x| *x & 1 != 0 );
+    /// assert_eq!(&array[..], &[1, 3]);
+    /// ```
+    pub fn retain<F>(&mut self, mut f: F)
+        where F: FnMut(&mut A::Item) -> bool
+    {
+        let len = self.len();
+        let mut del = 0;
+        {
+            let v = &mut **self;
+
+            for i in 0..len {
+                if !f(&mut v[i]) {
+                    del += 1;
+                } else if del > 0 {
+                    v.swap(i - del, i);
+                }
+            }
+        }
+        if del > 0 {
+            self.drain(len - del..);
+        }
+    }
+
     /// Set the vector's length without dropping or moving out elements
     ///
     /// May panic if `length` is greater than the capacity.
diff --git a/tests/tests.rs b/tests/tests.rs
index aa5b80d..aadacce 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -158,6 +158,23 @@ fn test_drain() {
     assert_eq!(&v[..], &[]);
 }
 
+#[test]
+fn test_retain() {
+    let mut v = ArrayVec::from([0; 8]);
+    for (i, elt) in v.iter_mut().enumerate() {
+        *elt = i;
+    }
+    v.retain(|_| true);
+    assert_eq!(&v[..], &[0, 1, 2, 3, 4, 5, 6, 7]);
+    v.retain(|elt| {
+        *elt /= 2;
+        *elt % 2 == 0
+    });
+    assert_eq!(&v[..], &[0, 0, 2, 2]);
+    v.retain(|_| false);
+    assert_eq!(&v[..], &[]);
+}
+
 #[test]
 #[should_panic]
 fn test_drain_oob() {