Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Vec::drain_filter #43245

Merged
merged 1 commit into from
Aug 16, 2017
Merged

Add Vec::drain_filter #43245

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/liballoc/tests/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#![feature(inclusive_range_syntax)]
#![feature(collection_placement)]
#![feature(const_fn)]
#![feature(drain_filter)]
#![feature(exact_size_is_empty)]
#![feature(iterator_step_by)]
#![feature(pattern)]
Expand Down
167 changes: 167 additions & 0 deletions src/liballoc/tests/vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -801,3 +801,170 @@ fn overaligned_allocations() {
assert!(v.as_ptr() as usize & 0xff == 0);
}
}

#[test]
fn drain_filter_empty() {
let mut vec: Vec<i32> = vec![];

{
let mut iter = vec.drain_filter(|_| true);
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
}
assert_eq!(vec.len(), 0);
assert_eq!(vec, vec![]);
}

#[test]
fn drain_filter_zst() {
let mut vec = vec![(), (), (), (), ()];
let initial_len = vec.len();
let mut count = 0;
{
let mut iter = vec.drain_filter(|_| true);
assert_eq!(iter.size_hint(), (0, Some(initial_len)));
while let Some(_) = iter.next() {
count += 1;
assert_eq!(iter.size_hint(), (0, Some(initial_len - count)));
}
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
}

assert_eq!(count, initial_len);
assert_eq!(vec.len(), 0);
assert_eq!(vec, vec![]);
}

#[test]
fn drain_filter_false() {
let mut vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];

let initial_len = vec.len();
let mut count = 0;
{
let mut iter = vec.drain_filter(|_| false);
assert_eq!(iter.size_hint(), (0, Some(initial_len)));
for _ in iter.by_ref() {
count += 1;
}
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
}

assert_eq!(count, 0);
assert_eq!(vec.len(), initial_len);
assert_eq!(vec, vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
}

#[test]
fn drain_filter_true() {
let mut vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];

let initial_len = vec.len();
let mut count = 0;
{
let mut iter = vec.drain_filter(|_| true);
assert_eq!(iter.size_hint(), (0, Some(initial_len)));
while let Some(_) = iter.next() {
count += 1;
assert_eq!(iter.size_hint(), (0, Some(initial_len - count)));
}
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
}

assert_eq!(count, initial_len);
assert_eq!(vec.len(), 0);
assert_eq!(vec, vec![]);
}

#[test]
fn drain_filter_complex() {

{ // [+xxx++++++xxxxx++++x+x++]
let mut vec = vec![1,
2, 4, 6,
7, 9, 11, 13, 15, 17,
18, 20, 22, 24, 26,
27, 29, 31, 33,
34,
35,
36,
37, 39];

let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
assert_eq!(removed.len(), 10);
assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);

assert_eq!(vec.len(), 14);
assert_eq!(vec, vec![1, 7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35, 37, 39]);
}

{ // [xxx++++++xxxxx++++x+x++]
let mut vec = vec![2, 4, 6,
7, 9, 11, 13, 15, 17,
18, 20, 22, 24, 26,
27, 29, 31, 33,
34,
35,
36,
37, 39];

let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
assert_eq!(removed.len(), 10);
assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);

assert_eq!(vec.len(), 13);
assert_eq!(vec, vec![7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35, 37, 39]);
}

{ // [xxx++++++xxxxx++++x+x]
let mut vec = vec![2, 4, 6,
7, 9, 11, 13, 15, 17,
18, 20, 22, 24, 26,
27, 29, 31, 33,
34,
35,
36];

let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
assert_eq!(removed.len(), 10);
assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);

assert_eq!(vec.len(), 11);
assert_eq!(vec, vec![7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35]);
}

{ // [xxxxxxxxxx+++++++++++]
let mut vec = vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
1, 3, 5, 7, 9, 11, 13, 15, 17, 19];

let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
assert_eq!(removed.len(), 10);
assert_eq!(removed, vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]);

assert_eq!(vec.len(), 10);
assert_eq!(vec, vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
}

{ // [+++++++++++xxxxxxxxxx]
let mut vec = vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19,
2, 4, 6, 8, 10, 12, 14, 16, 18, 20];

let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
assert_eq!(removed.len(), 10);
assert_eq!(removed, vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]);

assert_eq!(vec.len(), 10);
assert_eq!(vec, vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
}
}


113 changes: 113 additions & 0 deletions src/liballoc/vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1960,6 +1960,65 @@ impl<T> Vec<T> {
}
}

/// Creates an iterator which uses a closure to determine if an element should be removed.
///
/// If the closure returns true, then the element is removed and yielded.
/// If the closure returns false, it will try again, and call the closure
/// on the next element, seeing if it passes the test.
///
/// Using this method is equivalent to the following code:
///
/// ```
/// # let some_predicate = |x: &mut i32| { *x == 2 };
/// # let mut vec = vec![1, 2, 3, 4, 5];
/// let mut i = 0;
/// while i != vec.len() {
/// if some_predicate(&mut vec[i]) {
/// let val = vec.remove(i);
/// // your code here
/// }
/// i += 1;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would skip every item after a removed item.

let some_predicate = |&x| x == 2;
let mut vec = vec![1, 2, 2, 3, 4, 5]; 

let mut i = 0;
while i != vec.len() {
    if some_predicate(&mut vec[i]) {
        let val = vec.remove(i);
        println!("removed {}", val);
    }
    i += 1;
}

println!("{:?}", vec);

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oops, haha! Need to put the +=1 in an else

/// }
/// ```
///
/// But `drain_filter` is easier to use. `drain_filter` is also more efficient,
/// because it can backshift the elements of the array in bulk.
///
/// Note that `drain_filter` also lets you mutate every element in the filter closure,
/// regardless of whether you choose to keep or remove it.
///
///
/// # Examples
///
/// Splitting an array into evens and odds, reusing the original allocation:
///
/// ```
/// #![feature(drain_filter)]
/// let mut numbers = vec![1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15];
///
/// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
/// let odds = numbers;
///
/// assert_eq!(evens, vec![2, 4, 6, 8, 14]);
/// assert_eq!(odds, vec![1, 3, 5, 9, 11, 13, 15]);
/// ```
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
pub fn drain_filter<F>(&mut self, filter: F) -> DrainFilter<T, F>
where F: FnMut(&mut T) -> bool,
{
let old_len = self.len();

// Guard against us getting leaked (leak amplification)
unsafe { self.set_len(0); }

DrainFilter {
vec: self,
idx: 0,
del: 0,
old_len,
pred: filter,
}
}
}

#[stable(feature = "extend_ref", since = "1.2.0")]
Expand Down Expand Up @@ -2596,3 +2655,57 @@ impl<'a, T> Drain<'a, T> {
self.tail_start = new_tail_start;
}
}

/// An iterator produced by calling `drain_filter` on Vec.
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
#[derive(Debug)]
pub struct DrainFilter<'a, T: 'a, F>
where F: FnMut(&mut T) -> bool,
{
vec: &'a mut Vec<T>,
idx: usize,
del: usize,
old_len: usize,
pred: F,
}

#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
impl<'a, T, F> Iterator for DrainFilter<'a, T, F>
where F: FnMut(&mut T) -> bool,
{
type Item = T;

fn next(&mut self) -> Option<T> {
unsafe {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would personally not make the unsafe cover the entire block, just the raw_parts_mut method.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's also needed for the ptr::read. At that point might as well just say the whole thing is unsafe and call it a day (this matches many parts of this file).

while self.idx != self.old_len {
let i = self.idx;
self.idx += 1;
let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
if (self.pred)(&mut v[i]) {
self.del += 1;
return Some(ptr::read(&v[i]));
} else if self.del > 0 {
v.swap(i - self.del, i);
}
}
None
}
}

fn size_hint(&self) -> (usize, Option<usize>) {
(0, Some(self.old_len - self.idx))
}
}

#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
impl<'a, T, F> Drop for DrainFilter<'a, T, F>
where F: FnMut(&mut T) -> bool,
{
fn drop(&mut self) {
for _ in self.by_ref() { }

unsafe {
self.vec.set_len(self.old_len - self.del);
}
}
}