From ae7ff95f38817c13a632be15b1cef82359c0aa84 Mon Sep 17 00:00:00 2001
From: Louis Cloete <louisc363@gmail.com>
Date: Tue, 27 Dec 2022 00:15:05 +0200
Subject: [PATCH] DrainFiilter drain_keep_rest feature flag and test

---
 Cargo.toml   |  1 +
 src/lib.rs   | 26 ++++++++++++++++++++++++--
 src/tests.rs | 14 ++++++++++++++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index e2034ef..2963015 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,7 @@ union = []
 specialization = []
 may_dangle = []
 drain_filter = []
+drain_keep_rest = ["drain_filter"]
 
 # UNSTABLE FEATURES (requires Rust nightly)
 # Enable to use the #[debugger_visualizer] attribute.
diff --git a/src/lib.rs b/src/lib.rs
index be9226b..79282c3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -125,6 +125,9 @@ use core::marker::PhantomData;
 #[cfg(feature = "write")]
 use std::io;
 
+#[cfg(feature = "drain_keep_rest")]
+use core::mem::ManuallyDrop;
+
 /// Creates a [`SmallVec`] containing the arguments.
 ///
 /// `smallvec!` allows `SmallVec`s to be defined with the same syntax as array expressions.
@@ -541,12 +544,31 @@ where
     }
 }
 
-#[cfg(feature = "drain_filter")]
+#[cfg(feature = "drain_keep_rest")]
 impl <T, F> DrainFilter<'_, T, F>
 where
-    F: FnMut(T::Item) -> bool,
+    F: FnMut(&mut T::Item) -> bool,
     T: Array
 {
+    /// Keep unyielded elements in the source `Vec`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use smallvec::{smallvec, SmallVec};
+    ///
+    /// let mut vec: SmallVec<[char; 2]> = smallvec!['a', 'b', 'c'];
+    /// let mut drain = vec.drain_filter(|_| true);
+    ///
+    /// assert_eq!(drain.next().unwrap(), 'a');
+    ///
+    /// // This call keeps 'b' and 'c' in the vec.
+    /// drain.keep_rest();
+    ///
+    /// // If we wouldn't call `keep_rest()`,
+    /// // `vec` would be empty.
+    /// assert_eq!(vec, SmallVec::<[char; 2]>::from_slice(&['b', 'c']));
+    /// ```
     pub fn keep_rest(self)
     {
         // At this moment layout looks like this:
diff --git a/src/tests.rs b/src/tests.rs
index 9627615..bb39dde 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1000,3 +1000,17 @@ fn drain_filter() {
     assert_eq!(a, SmallVec::<[u8; 2]>::from_slice(&[1u8, 2, 4, 5, 7, 8]));
     assert_eq!(b, SmallVec::<[u8; 2]>::from_slice(&[3u8, 6]));
 }
+
+#[cfg(feature = "drain_keep_rest")]
+#[test]
+fn drain_keep_rest() {
+    let mut a: SmallVec<[i32; 3]> = smallvec![1i32, 2, 3, 4, 5, 6, 7, 8];
+    let mut df = a.drain_filter(|x| *x % 2 == 0);
+
+    assert_eq!(df.next().unwrap(), 2);
+    assert_eq!(df.next().unwrap(), 4);
+
+    df.keep_rest();
+
+    assert_eq!(a, SmallVec::<[i32; 3]>::from_slice(&[1i32, 3, 5, 6, 7, 8]));
+}