1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

use std::ptr;

/// Trait that provides `retain_mut` method.
pub trait RetainMut<T> {
    /// Retains only the elements specified by the predicate with a mutable borrow.
    ///
    /// In other words, remove all elements `e` such that `f(&mut e)` returns `false`.
    /// This method operates in place and preserves the order of the retained
    /// elements.
    ///
    /// # Examples
    ///
    /// ```
    /// use retain_mut::RetainMut;
    ///
    /// let mut vec = vec![1, 2, 3, 4];
    /// vec.retain_mut(|x| { *x *= 3; *x % 2 == 0 });
    /// assert_eq!(vec, [6, 12]);
    /// ```
    fn retain_mut<F>(&mut self, f: F) where F: FnMut(&mut T) -> bool;
}

impl<T> RetainMut<T> for Vec<T> {
    // The implementation is based on
    // https://github.com/rust-lang/rust/blob/a67749ae87b1c873ed09fca2a204beff2fe5e7ea/src/liballoc/vec.rs#L804-L829
    fn retain_mut<F>(&mut self, mut f: F)
    where
        F: FnMut(&mut T) -> bool,
    {
        let len = self.len();
        let mut del = 0;
        {
            let v = &mut **self;
            for i in 0..len {
                if !f(&mut v[i]) {
                    del += 1;
                    unsafe {
                        ptr::drop_in_place(&mut v[i]);
                    }
                } else if del > 0 {
                    let src: *const T = &v[i];
                    let dst: *mut T = &mut v[i - del];
                    unsafe {
                        ptr::copy_nonoverlapping(src, dst, 1);
                    }
                }
            }
        }
        unsafe {
            self.set_len(len - del);
        }
    }
}