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}