cs_utils/utils/
drain_filter.rs

1
2/// Drain and filter at the same time, accepts a `mutable vector ref` and `predicate`.
3/// The `predicate` receives a `vector element ref` for the current iteration and
4/// in-place removes the element from the vector, if `predicate` returns `true`.
5/// Returns and array of `removed` elements back.
6/// 
7/// ## Examples
8/// 
9/// ```
10/// use cs_utils::{drain_filter, random_bool};
11/// 
12/// let mut test_vector = vec![
13///     random_bool(),
14///     random_bool(),
15///     random_bool(),
16///     random_bool(),
17///     random_bool(),
18///     random_bool(),
19///     random_bool(),
20///     random_bool(),
21///     random_bool(),
22///     random_bool(),
23///     random_bool(),
24/// ];
25/// 
26/// let removed_items = drain_filter(
27///     &mut test_vector,
28///     |el| {
29///         return *el;
30///     },
31/// );
32/// 
33/// assert!(
34///     !removed_items.contains(&false),
35///     "Removed items must not contain \"false\" values.",
36/// );
37/// 
38/// assert!(
39///     !test_vector.contains(&true),
40///     "Removed items must not contain \"true\" values.",
41/// );
42/// ```
43pub fn drain_filter<T>(vector: &mut Vec<T>, mut predicate: impl FnMut(&T) -> bool) -> Vec<T> {
44    let mut result = vec![];
45
46    let mut i = 0;
47        while i < vector.len() {
48            let item = &vector[i];
49            if predicate(item) {
50                let removed_item = vector.swap_remove(i);
51                result.push(removed_item);
52                // since we swapped with the last vector item, don't increment
53                // the index, - the swapped item must be checked too
54                continue;
55            }
56
57            i += 1;
58        }
59
60    return result;
61}
62
63#[cfg(test)]
64mod tests {
65    use crate::drain_filter;
66
67    #[test]
68    fn drains_the_vector() {
69        let mut test_vector = vec![
70            true,
71            false,
72            false,
73            true,
74            true,
75            false,
76            false,
77            false,
78        ];
79
80        let removed = drain_filter(&mut test_vector, |el| {
81            return *el;
82        });
83
84        assert_eq!(
85            removed.len(),
86            3,
87            "Must remove all elements equal to \"true\"."
88        );
89
90        assert!(
91            !removed.contains(&false),
92            "Removed elements vector must not contain \"false\" values."
93        );
94
95        assert_eq!(
96            test_vector.len(),
97            5,
98            "Original vector must retain \"false\" values."
99        );
100
101        assert!(
102            !test_vector.contains(&true),
103            "Original vector must not contain \"true\" values."
104        );
105    }
106
107    #[test]
108    fn works_with_non_clone_items() {
109        struct TestStruct {
110            pub flag: bool,
111        }
112
113        let mut test_vector = vec![
114            TestStruct { flag: false },
115            TestStruct { flag: true },
116            TestStruct { flag: false },
117            TestStruct { flag: true },
118            TestStruct { flag: false },
119            TestStruct { flag: true },
120            TestStruct { flag: false },
121            TestStruct { flag: true },
122            TestStruct { flag: true },
123        ];
124
125        let removed = drain_filter(&mut test_vector, |el| {
126            return el.flag;
127        });
128
129        assert_eq!(
130            removed.len(),
131            5,
132            "Must remove all elements equal to \"true\"."
133        );
134
135        assert_eq!(
136            test_vector.len(),
137            4,
138            "Must retain all elements equal to \"false\"."
139        );
140    }
141}