Skip to content

Commit

Permalink
DrainFiilter drain_keep_rest feature flag and test
Browse files Browse the repository at this point in the history
  • Loading branch information
L0uisc authored and mbrubeck committed Jun 15, 2023
1 parent 5d0c7e2 commit ae7ff95
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 2 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
26 changes: 24 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
14 changes: 14 additions & 0 deletions src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]));
}

0 comments on commit ae7ff95

Please sign in to comment.