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 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])); +}