[Pre-RFC] adding retain_or_drain method to Vec<T> that allows you retrieve the elements that are deleted from retain

Hi,

A couple of times now, I have wanted to iterate through a vector, delete some elements that I now consider to be inactive, and then do an action for each of the deleted elements. So far I have had to do something similar to this:

let mut vec = vec![1, 2, 3, 4];
vec.retain(|x| if x % 2 == 0 {
  true
} else {
  // use the element that will be deleted to do something
  println!("elem not to be retained: {:?}", x);
  false
});
println!("remaining of vector: {:?}", vec);

This is fairly unelegant and hard to read in my opinion. We have had to combine a side-effect inside what should be a clean retain check.

The other ways to delete an element from a vector (swap_remove and remove) return the deleted element, so it would be consistent to have a method like retain that also returns the deleted elements. Since vector already implements drain, my idea would be to have a method that retains objects if they match a predicate (just as the retain now) or returns them inside a Drain if it doesn’t. Something like this:

fn retain_or_drain<F>(&mut self, mut f: F) -> std::vec::Drain<T>
  where
    F: FnMut(&T) -> bool,
{
  // this is exactly the same as the current retain() method...
  let len = self.len();
  let mut del = 0;
  {
    let v = &mut **self;

    for i in 0..len {
      if !f(&v[i]) {
        del += 1;
      } else if del > 0 {
        v.swap(i - del, i);
      }
    }
  }
  //... until here
  // the existing retain() methods truncates
  // this new method would return a drain instead
  self.drain(len - del..)
}

This method would allow the example above to be rewritten as:

let mut vec = vec![1, 2, 3, 4];
for k in vec.retain_or_drain(|x| x % 2 == 0) {
  println!("drained element: {:?}", k);
}
println!("remaining of vector: {:?}", vec);

This is easier to read and understand. I first do the retain check, and I am returned a drain with all the elements that were not retained. I can now do the side-effecty things I wanted to do with those elements, without combining that logic inside my retain check.

This is my first time suggesting an RFC so I wanted to get the opinions from you all before I decided to move forward towards creating an official RFC. Is there a better way to delete all elements that match a criteria, and do something with each of these elements without introducing more features to the std vector?. I am open to any suggestions on this idea as well as how to better create a Pre-RFC or RFC in the future.

A related method that I would like is partition, which moves elements of a slice to the left or right end based on a predicate . This is available in the C++ standard library as std::partition, and it is also available in itertools. A simple version for slices would look like this:

Update: This implementation is incorrect, because it does not preserve the order of elements that end up at the back of the vector.

fn partition<T, F>(v: &mut [T], mut f: F) -> usize
where F: FnMut(&T) -> bool
{
    let len = v.len();
    let mut del = 0;

    for i in 0..len {
        if !f(&v[i]) {
            del += 1;
        } else if del > 0 {
            v.swap(i - del, i);
        }
    }
    len - del
}

It is a useful building block for methods like retain. For example, Vec::<T>::retain_or_drain could then be implemented as:

fn retain_or_drain<F>(&mut self, mut f: F) -> std::vec::Drain<T>
where F: FnMut(&T) -> bool
{
    let i = partition(self, f);
    self.drain(i..)
}
2 Likes

This will be called Vec::drain_filter (PR: 43245) and hopefully will land on nightly soon.

2 Likes

This topic was automatically closed 90 days after the last reply. New replies are no longer allowed.