Skip to content

Commit 0d244f4

Browse files
committed
port feature drain_filter (servo#292) to v2
1 parent a410c67 commit 0d244f4

File tree

4 files changed

+265
-0
lines changed

4 files changed

+265
-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
@@ -16,6 +16,8 @@ documentation = "https://docs.rs/smallvec/"
1616
write = []
1717
specialization = []
1818
may_dangle = []
19+
drain_filter = []
20+
drain_keep_rest = ["drain_filter"]
1921

2022
[dependencies]
2123
serde = { version = "1", optional = true, default-features = false }

src/lib.rs

+237
Original file line numberDiff line numberDiff line change
@@ -380,6 +380,178 @@ impl<'a, T: 'a, const N: usize> Drop for Drain<'a, T, N> {
380380
}
381381
}
382382

383+
#[cfg(feature = "drain_filter")]
384+
/// An iterator which uses a closure to determine if an element should be removed.
385+
///
386+
/// Returned from [`SmallVec::drain_filter`][1].
387+
///
388+
/// [1]: struct.SmallVec.html#method.drain_filter
389+
pub struct DrainFilter<'a, T, const N: usize, F>
390+
where
391+
F: FnMut(&mut T) -> bool,
392+
{
393+
vec: &'a mut SmallVec<T, N>,
394+
/// The index of the item that will be inspected by the next call to `next`.
395+
idx: usize,
396+
/// The number of items that have been drained (removed) thus far.
397+
del: usize,
398+
/// The original length of `vec` prior to draining.
399+
old_len: usize,
400+
/// The filter test predicate.
401+
pred: F,
402+
/// A flag that indicates a panic has occurred in the filter test predicate.
403+
/// This is used as a hint in the drop implementation to prevent consumption
404+
/// of the remainder of the `DrainFilter`. Any unprocessed items will be
405+
/// backshifted in the `vec`, but no further items will be dropped or
406+
/// tested by the filter predicate.
407+
panic_flag: bool,
408+
}
409+
410+
#[cfg(feature = "drain_filter")]
411+
impl<T, F, const N: usize> Iterator for DrainFilter<'_, T, N, F>
412+
where
413+
F: FnMut(&mut T) -> bool,
414+
{
415+
type Item = T;
416+
417+
fn next(&mut self) -> Option<T> {
418+
unsafe {
419+
while self.idx < self.old_len {
420+
let i = self.idx;
421+
let v = core::slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
422+
self.panic_flag = true;
423+
let drained = (self.pred)(&mut v[i]);
424+
self.panic_flag = false;
425+
// Update the index *after* the predicate is called. If the index
426+
// is updated prior and the predicate panics, the element at this
427+
// index would be leaked.
428+
self.idx += 1;
429+
if drained {
430+
self.del += 1;
431+
return Some(core::ptr::read(&v[i]));
432+
} else if self.del > 0 {
433+
let del = self.del;
434+
let src: *const Self::Item = &v[i];
435+
let dst: *mut Self::Item = &mut v[i - del];
436+
core::ptr::copy_nonoverlapping(src, dst, 1);
437+
}
438+
}
439+
None
440+
}
441+
}
442+
443+
fn size_hint(&self) -> (usize, Option<usize>) {
444+
(0, Some(self.old_len - self.idx))
445+
}
446+
}
447+
448+
#[cfg(feature = "drain_filter")]
449+
impl<T, F, const N: usize> Drop for DrainFilter<'_, T, N, F>
450+
where
451+
F: FnMut(&mut T) -> bool,
452+
{
453+
fn drop(&mut self) {
454+
struct BackshiftOnDrop<'a, 'b, T, const N: usize, F>
455+
where
456+
F: FnMut(&mut T) -> bool,
457+
{
458+
drain: &'b mut DrainFilter<'a, T, N, F>,
459+
}
460+
461+
impl<'a, 'b, T, const N: usize, F> Drop for BackshiftOnDrop<'a, 'b, T, N, F>
462+
where
463+
F: FnMut(&mut T) -> bool,
464+
{
465+
fn drop(&mut self) {
466+
unsafe {
467+
if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
468+
// This is a pretty messed up state, and there isn't really an
469+
// obviously right thing to do. We don't want to keep trying
470+
// to execute `pred`, so we just backshift all the unprocessed
471+
// elements and tell the vec that they still exist. The backshift
472+
// is required to prevent a double-drop of the last successfully
473+
// drained item prior to a panic in the predicate.
474+
let ptr = self.drain.vec.as_mut_ptr();
475+
let src = ptr.add(self.drain.idx);
476+
let dst = src.sub(self.drain.del);
477+
let tail_len = self.drain.old_len - self.drain.idx;
478+
src.copy_to(dst, tail_len);
479+
}
480+
self.drain.vec.set_len(self.drain.old_len - self.drain.del);
481+
}
482+
}
483+
}
484+
485+
let backshift = BackshiftOnDrop { drain: self };
486+
487+
// Attempt to consume any remaining elements if the filter predicate
488+
// has not yet panicked. We'll backshift any remaining elements
489+
// whether we've already panicked or if the consumption here panics.
490+
if !backshift.drain.panic_flag {
491+
backshift.drain.for_each(drop);
492+
}
493+
}
494+
}
495+
496+
#[cfg(feature = "drain_keep_rest")]
497+
impl<T, F, const N: usize> DrainFilter<'_, T, N, F>
498+
where
499+
F: FnMut(&mut T) -> bool,
500+
{
501+
/// Keep unyielded elements in the source `Vec`.
502+
///
503+
/// # Examples
504+
///
505+
/// ```
506+
/// # use smallvec::{smallvec, SmallVec};
507+
///
508+
/// let mut vec: SmallVec<char, 2> = smallvec!['a', 'b', 'c'];
509+
/// let mut drain = vec.drain_filter(|_| true);
510+
///
511+
/// assert_eq!(drain.next().unwrap(), 'a');
512+
///
513+
/// // This call keeps 'b' and 'c' in the vec.
514+
/// drain.keep_rest();
515+
///
516+
/// // If we wouldn't call `keep_rest()`,
517+
/// // `vec` would be empty.
518+
/// assert_eq!(vec, SmallVec::<char, 2>::from_slice(&['b', 'c']));
519+
/// ```
520+
pub fn keep_rest(self) {
521+
// At this moment layout looks like this:
522+
//
523+
// _____________________/-- old_len
524+
// / \
525+
// [kept] [yielded] [tail]
526+
// \_______/ ^-- idx
527+
// \-- del
528+
//
529+
// Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
530+
//
531+
// 1. Move [tail] after [kept]
532+
// 2. Update length of the original vec to `old_len - del`
533+
// a. In case of ZST, this is the only thing we want to do
534+
// 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
535+
let mut this = ManuallyDrop::new(self);
536+
537+
unsafe {
538+
// ZSTs have no identity, so we don't need to move them around.
539+
let needs_move = core::mem::size_of::<T>() != 0;
540+
541+
if needs_move && this.idx < this.old_len && this.del > 0 {
542+
let ptr = this.vec.as_mut_ptr();
543+
let src = ptr.add(this.idx);
544+
let dst = src.sub(this.del);
545+
let tail_len = this.old_len - this.idx;
546+
src.copy_to(dst, tail_len);
547+
}
548+
549+
let new_len = this.old_len - this.del;
550+
this.vec.set_len(new_len);
551+
}
552+
}
553+
}
554+
383555
/// An iterator that consumes a `SmallVec` and yields its items by value.
384556
///
385557
/// Returned from [`SmallVec::into_iter`][1].
@@ -732,6 +904,71 @@ impl<T, const N: usize> SmallVec<T, N> {
732904
}
733905
}
734906

907+
#[cfg(feature = "drain_filter")]
908+
/// Creates an iterator which uses a closure to determine if an element should be removed.
909+
///
910+
/// If the closure returns true, the element is removed and yielded. If the closure returns
911+
/// false, the element will remain in the vector and will not be yielded by the iterator.
912+
///
913+
/// Using this method is equivalent to the following code:
914+
/// ```
915+
/// # use smallvec::SmallVec;
916+
/// # let some_predicate = |x: &mut i32| { *x == 2 || *x == 3 || *x == 6 };
917+
/// # let mut vec: SmallVec<i32, 8> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6]);
918+
/// let mut i = 0;
919+
/// while i < vec.len() {
920+
/// if some_predicate(&mut vec[i]) {
921+
/// let val = vec.remove(i);
922+
/// // your code here
923+
/// } else {
924+
/// i += 1;
925+
/// }
926+
/// }
927+
///
928+
/// # assert_eq!(vec, SmallVec::<i32, 8>::from_slice(&[1i32, 4, 5]));
929+
/// ```
930+
/// ///
931+
/// But `drain_filter` is easier to use. `drain_filter` is also more efficient,
932+
/// because it can backshift the elements of the array in bulk.
933+
///
934+
/// Note that `drain_filter` also lets you mutate every element in the filter closure,
935+
/// regardless of whether you choose to keep or remove it.
936+
///
937+
/// # Examples
938+
///
939+
/// Splitting an array into evens and odds, reusing the original allocation:
940+
///
941+
/// ```
942+
/// # use smallvec::SmallVec;
943+
/// let mut numbers: SmallVec<i32, 16> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
944+
///
945+
/// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<SmallVec<i32, 16>>();
946+
/// let odds = numbers;
947+
///
948+
/// assert_eq!(evens, SmallVec::<i32, 16>::from_slice(&[2i32, 4, 6, 8, 14]));
949+
/// assert_eq!(odds, SmallVec::<i32, 16>::from_slice(&[1i32, 3, 5, 9, 11, 13, 15]));
950+
/// ```
951+
pub fn drain_filter<F>(&mut self, filter: F) -> DrainFilter<'_, T, N, F>
952+
where
953+
F: FnMut(&mut T) -> bool,
954+
{
955+
let old_len = self.len();
956+
957+
// Guard against us getting leaked (leak amplification)
958+
unsafe {
959+
self.set_len(0);
960+
}
961+
962+
DrainFilter {
963+
vec: self,
964+
idx: 0,
965+
del: 0,
966+
old_len,
967+
pred: filter,
968+
panic_flag: false,
969+
}
970+
}
971+
735972
#[inline]
736973
pub fn push(&mut self, value: T) {
737974
let len = self.len();

src/tests.rs

+25
Original file line numberDiff line numberDiff line change
@@ -1003,3 +1003,28 @@ fn test_clone_from() {
10031003
b.clone_from(&c);
10041004
assert_eq!(&*b, &[20, 21, 22]);
10051005
}
1006+
1007+
#[cfg(feature = "drain_filter")]
1008+
#[test]
1009+
fn drain_filter() {
1010+
let mut a: SmallVec<u8, 2> = smallvec![1u8, 2, 3, 4, 5, 6, 7, 8];
1011+
1012+
let b: SmallVec<u8, 2> = a.drain_filter(|x| *x % 3 == 0).collect();
1013+
1014+
assert_eq!(a, SmallVec::<u8, 2>::from_slice(&[1u8, 2, 4, 5, 7, 8]));
1015+
assert_eq!(b, SmallVec::<u8, 2>::from_slice(&[3u8, 6]));
1016+
}
1017+
1018+
#[cfg(feature = "drain_keep_rest")]
1019+
#[test]
1020+
fn drain_keep_rest() {
1021+
let mut a: SmallVec<i32, 3> = smallvec![1i32, 2, 3, 4, 5, 6, 7, 8];
1022+
let mut df = a.drain_filter(|x| *x % 2 == 0);
1023+
1024+
assert_eq!(df.next().unwrap(), 2);
1025+
assert_eq!(df.next().unwrap(), 4);
1026+
1027+
df.keep_rest();
1028+
1029+
assert_eq!(a, SmallVec::<i32, 3>::from_slice(&[1i32, 3, 5, 6, 7, 8]));
1030+
}

0 commit comments

Comments
 (0)