Skip to content

Commit 3e20520

Browse files
authored
Auto merge of #292 - L0uisc:drain_filter, r=mbrubeck
Implemented DrainFilter Provides an implementation for DrainFilter (issue #281).
2 parents 0c56f1b + ae7ff95 commit 3e20520

File tree

4 files changed

+282
-0
lines changed

4 files changed

+282
-0
lines changed

.gitignore

+1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
target
22
/Cargo.lock
33
/fuzz/hfuzz_target
4+
/.vscode

Cargo.toml

+2
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ write = []
1818
union = []
1919
specialization = []
2020
may_dangle = []
21+
drain_filter = []
22+
drain_keep_rest = ["drain_filter"]
2123

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

src/lib.rs

+254
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.
@@ -410,6 +413,198 @@ impl<'a, T: 'a + Array> Drop for Drain<'a, T> {
410413
}
411414
}
412415

416+
#[cfg(feature = "drain_filter")]
417+
/// An iterator which uses a closure to determine if an element should be removed.
418+
///
419+
/// Returned from [`SmallVec::drain_filter`][1].
420+
///
421+
/// [1]: struct.SmallVec.html#method.drain_filter
422+
pub struct DrainFilter<'a, T, F>
423+
where
424+
F: FnMut(&mut T::Item) -> bool,
425+
T: Array,
426+
{
427+
vec: &'a mut SmallVec<T>,
428+
/// The index of the item that will be inspected by the next call to `next`.
429+
idx: usize,
430+
/// The number of items that have been drained (removed) thus far.
431+
del: usize,
432+
/// The original length of `vec` prior to draining.
433+
old_len: usize,
434+
/// The filter test predicate.
435+
pred: F,
436+
/// A flag that indicates a panic has occurred in the filter test predicate.
437+
/// This is used as a hint in the drop implementation to prevent consumption
438+
/// of the remainder of the `DrainFilter`. Any unprocessed items will be
439+
/// backshifted in the `vec`, but no further items will be dropped or
440+
/// tested by the filter predicate.
441+
panic_flag: bool,
442+
}
443+
444+
#[cfg(feature = "drain_filter")]
445+
impl <T, F> fmt::Debug for DrainFilter<'_, T, F>
446+
where
447+
F: FnMut(&mut T::Item) -> bool,
448+
T: Array,
449+
T::Item: fmt::Debug,
450+
{
451+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
452+
f.debug_tuple("DrainFilter").field(&self.vec.as_slice()).finish()
453+
}
454+
}
455+
456+
#[cfg(feature = "drain_filter")]
457+
impl <T, F> Iterator for DrainFilter<'_, T, F>
458+
where
459+
F: FnMut(&mut T::Item) -> bool,
460+
T: Array,
461+
{
462+
type Item = T::Item;
463+
464+
fn next(&mut self) -> Option<T::Item>
465+
{
466+
unsafe {
467+
while self.idx < self.old_len {
468+
let i = self.idx;
469+
let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
470+
self.panic_flag = true;
471+
let drained = (self.pred)(&mut v[i]);
472+
self.panic_flag = false;
473+
// Update the index *after* the predicate is called. If the index
474+
// is updated prior and the predicate panics, the element at this
475+
// index would be leaked.
476+
self.idx += 1;
477+
if drained {
478+
self.del += 1;
479+
return Some(ptr::read(&v[i]));
480+
} else if self.del > 0 {
481+
let del = self.del;
482+
let src: *const Self::Item = &v[i];
483+
let dst: *mut Self::Item = &mut v[i - del];
484+
ptr::copy_nonoverlapping(src, dst, 1);
485+
}
486+
}
487+
None
488+
}
489+
}
490+
491+
fn size_hint(&self) -> (usize, Option<usize>) {
492+
(0, Some(self.old_len - self.idx))
493+
}
494+
}
495+
496+
#[cfg(feature = "drain_filter")]
497+
impl <T, F> Drop for DrainFilter<'_, T, F>
498+
where
499+
F: FnMut(&mut T::Item) -> bool,
500+
T: Array,
501+
{
502+
fn drop(&mut self) {
503+
struct BackshiftOnDrop<'a, 'b, T, F>
504+
where
505+
F: FnMut(&mut T::Item) -> bool,
506+
T: Array
507+
{
508+
drain: &'b mut DrainFilter<'a, T, F>,
509+
}
510+
511+
impl<'a, 'b, T, F> Drop for BackshiftOnDrop<'a, 'b, T, F>
512+
where
513+
F: FnMut(&mut T::Item) -> bool,
514+
T: Array
515+
{
516+
fn drop(&mut self) {
517+
unsafe {
518+
if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
519+
// This is a pretty messed up state, and there isn't really an
520+
// obviously right thing to do. We don't want to keep trying
521+
// to execute `pred`, so we just backshift all the unprocessed
522+
// elements and tell the vec that they still exist. The backshift
523+
// is required to prevent a double-drop of the last successfully
524+
// drained item prior to a panic in the predicate.
525+
let ptr = self.drain.vec.as_mut_ptr();
526+
let src = ptr.add(self.drain.idx);
527+
let dst = src.sub(self.drain.del);
528+
let tail_len = self.drain.old_len - self.drain.idx;
529+
src.copy_to(dst, tail_len);
530+
}
531+
self.drain.vec.set_len(self.drain.old_len - self.drain.del);
532+
}
533+
}
534+
}
535+
536+
let backshift = BackshiftOnDrop { drain: self };
537+
538+
// Attempt to consume any remaining elements if the filter predicate
539+
// has not yet panicked. We'll backshift any remaining elements
540+
// whether we've already panicked or if the consumption here panics.
541+
if !backshift.drain.panic_flag {
542+
backshift.drain.for_each(drop);
543+
}
544+
}
545+
}
546+
547+
#[cfg(feature = "drain_keep_rest")]
548+
impl <T, F> DrainFilter<'_, T, F>
549+
where
550+
F: FnMut(&mut T::Item) -> bool,
551+
T: Array
552+
{
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+
/// ```
572+
pub fn keep_rest(self)
573+
{
574+
// At this moment layout looks like this:
575+
//
576+
// _____________________/-- old_len
577+
// / \
578+
// [kept] [yielded] [tail]
579+
// \_______/ ^-- idx
580+
// \-- del
581+
//
582+
// Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
583+
//
584+
// 1. Move [tail] after [kept]
585+
// 2. Update length of the original vec to `old_len - del`
586+
// a. In case of ZST, this is the only thing we want to do
587+
// 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
588+
let mut this = ManuallyDrop::new(self);
589+
590+
unsafe {
591+
// ZSTs have no identity, so we don't need to move them around.
592+
let needs_move = mem::size_of::<T>() != 0;
593+
594+
if needs_move && this.idx < this.old_len && this.del > 0 {
595+
let ptr = this.vec.as_mut_ptr();
596+
let src = ptr.add(this.idx);
597+
let dst = src.sub(this.del);
598+
let tail_len = this.old_len - this.idx;
599+
src.copy_to(dst, tail_len);
600+
}
601+
602+
let new_len = this.old_len - this.del;
603+
this.vec.set_len(new_len);
604+
}
605+
}
606+
}
607+
413608
#[cfg(feature = "union")]
414609
union SmallVecData<A: Array> {
415610
inline: core::mem::ManuallyDrop<MaybeUninit<A>>,
@@ -847,6 +1042,65 @@ impl<A: Array> SmallVec<A> {
8471042
}
8481043
}
8491044

1045+
1046+
#[cfg(feature = "drain_filter")]
1047+
/// Creates an iterator which uses a closure to determine if an element should be removed.
1048+
///
1049+
/// If the closure returns true, the element is removed and yielded. If the closure returns
1050+
/// false, the element will remain in the vector and will not be yielded by the iterator.
1051+
///
1052+
/// Using this method is equivalent to the following code:
1053+
/// ```
1054+
/// # use smallvec::SmallVec;
1055+
/// # let some_predicate = |x: &mut i32| { *x == 2 || *x == 3 || *x == 6 };
1056+
/// # let mut vec: SmallVec<[i32; 8]> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6]);
1057+
/// let mut i = 0;
1058+
/// while i < vec.len() {
1059+
/// if some_predicate(&mut vec[i]) {
1060+
/// let val = vec.remove(i);
1061+
/// // your code here
1062+
/// } else {
1063+
/// i += 1;
1064+
/// }
1065+
/// }
1066+
///
1067+
/// # assert_eq!(vec, SmallVec::<[i32; 8]>::from_slice(&[1i32, 4, 5]));
1068+
/// ```
1069+
/// ///
1070+
/// But `drain_filter` is easier to use. `drain_filter` is also more efficient,
1071+
/// because it can backshift the elements of the array in bulk.
1072+
///
1073+
/// Note that `drain_filter` also lets you mutate every element in the filter closure,
1074+
/// regardless of whether you choose to keep or remove it.
1075+
///
1076+
/// # Examples
1077+
///
1078+
/// Splitting an array into evens and odds, reusing the original allocation:
1079+
///
1080+
/// ```
1081+
/// # use smallvec::SmallVec;
1082+
/// let mut numbers: SmallVec<[i32; 16]> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
1083+
///
1084+
/// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<SmallVec<[i32; 16]>>();
1085+
/// let odds = numbers;
1086+
///
1087+
/// assert_eq!(evens, SmallVec::<[i32; 16]>::from_slice(&[2i32, 4, 6, 8, 14]));
1088+
/// assert_eq!(odds, SmallVec::<[i32; 16]>::from_slice(&[1i32, 3, 5, 9, 11, 13, 15]));
1089+
/// ```
1090+
pub fn drain_filter<F>(&mut self, filter: F) -> DrainFilter<'_, A, F,>
1091+
where
1092+
F: FnMut(&mut A::Item) -> bool,
1093+
{
1094+
let old_len = self.len();
1095+
1096+
// Guard against us getting leaked (leak amplification)
1097+
unsafe {
1098+
self.set_len(0);
1099+
}
1100+
1101+
DrainFilter { vec: self, idx: 0, del: 0, old_len, pred: filter, panic_flag: false }
1102+
}
1103+
8501104
/// Append an item to the vector.
8511105
#[inline]
8521106
pub fn push(&mut self, value: A::Item) {

src/tests.rs

+25
Original file line numberDiff line numberDiff line change
@@ -989,3 +989,28 @@ fn test_size() {
989989
use core::mem::size_of;
990990
assert_eq!(24, size_of::<SmallVec<[u8; 8]>>());
991991
}
992+
993+
#[cfg(feature = "drain_filter")]
994+
#[test]
995+
fn drain_filter() {
996+
let mut a: SmallVec<[u8; 2]> = smallvec![1u8, 2, 3, 4, 5, 6, 7, 8];
997+
998+
let b: SmallVec<[u8; 2]> = a.drain_filter(|x| *x % 3 == 0).collect();
999+
1000+
assert_eq!(a, SmallVec::<[u8; 2]>::from_slice(&[1u8, 2, 4, 5, 7, 8]));
1001+
assert_eq!(b, SmallVec::<[u8; 2]>::from_slice(&[3u8, 6]));
1002+
}
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)