Skip to content

Commit ae7ff95

Browse files
L0uiscmbrubeck
authored andcommitted
DrainFiilter drain_keep_rest feature flag and test
1 parent 5d0c7e2 commit ae7ff95

File tree

3 files changed

+39
-2
lines changed

3 files changed

+39
-2
lines changed

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ union = []
1919
specialization = []
2020
may_dangle = []
2121
drain_filter = []
22+
drain_keep_rest = ["drain_filter"]
2223

2324
# UNSTABLE FEATURES (requires Rust nightly)
2425
# Enable to use the #[debugger_visualizer] attribute.

src/lib.rs

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,9 @@ use core::marker::PhantomData;
125125
#[cfg(feature = "write")]
126126
use std::io;
127127

128+
#[cfg(feature = "drain_keep_rest")]
129+
use core::mem::ManuallyDrop;
130+
128131
/// Creates a [`SmallVec`] containing the arguments.
129132
///
130133
/// `smallvec!` allows `SmallVec`s to be defined with the same syntax as array expressions.
@@ -541,12 +544,31 @@ where
541544
}
542545
}
543546

544-
#[cfg(feature = "drain_filter")]
547+
#[cfg(feature = "drain_keep_rest")]
545548
impl <T, F> DrainFilter<'_, T, F>
546549
where
547-
F: FnMut(T::Item) -> bool,
550+
F: FnMut(&mut T::Item) -> bool,
548551
T: Array
549552
{
553+
/// Keep unyielded elements in the source `Vec`.
554+
///
555+
/// # Examples
556+
///
557+
/// ```
558+
/// # use smallvec::{smallvec, SmallVec};
559+
///
560+
/// let mut vec: SmallVec<[char; 2]> = smallvec!['a', 'b', 'c'];
561+
/// let mut drain = vec.drain_filter(|_| true);
562+
///
563+
/// assert_eq!(drain.next().unwrap(), 'a');
564+
///
565+
/// // This call keeps 'b' and 'c' in the vec.
566+
/// drain.keep_rest();
567+
///
568+
/// // If we wouldn't call `keep_rest()`,
569+
/// // `vec` would be empty.
570+
/// assert_eq!(vec, SmallVec::<[char; 2]>::from_slice(&['b', 'c']));
571+
/// ```
550572
pub fn keep_rest(self)
551573
{
552574
// At this moment layout looks like this:

src/tests.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1000,3 +1000,17 @@ fn drain_filter() {
10001000
assert_eq!(a, SmallVec::<[u8; 2]>::from_slice(&[1u8, 2, 4, 5, 7, 8]));
10011001
assert_eq!(b, SmallVec::<[u8; 2]>::from_slice(&[3u8, 6]));
10021002
}
1003+
1004+
#[cfg(feature = "drain_keep_rest")]
1005+
#[test]
1006+
fn drain_keep_rest() {
1007+
let mut a: SmallVec<[i32; 3]> = smallvec![1i32, 2, 3, 4, 5, 6, 7, 8];
1008+
let mut df = a.drain_filter(|x| *x % 2 == 0);
1009+
1010+
assert_eq!(df.next().unwrap(), 2);
1011+
assert_eq!(df.next().unwrap(), 4);
1012+
1013+
df.keep_rest();
1014+
1015+
assert_eq!(a, SmallVec::<[i32; 3]>::from_slice(&[1i32, 3, 5, 6, 7, 8]));
1016+
}

0 commit comments

Comments
 (0)