retain_mut/
lib.rs

1//! **This crate has been deprecated.
2//! Rust 1.61 stabilized `retain_mut` for `Vec` and `VecDeque`,
3//! so you can use them directly.
4//! This crate is no longer maintained.**
5//!
6//! This crate provides trait `RetainMut` which
7//! provides `retain_mut` method for `Vec` and `VecDeque`.
8//!
9//! `retain_mut` is basically the same as `retain` except that
10//! it gives mutable reference of items to the predicate function.
11//!
12//! Since there is no reason `retain` couldn't have been designed this way,
13//! this crate basically just copies the code from std with minor changes
14//! to hand out mutable reference.
15//! The code these impls are based on can be found in code comments of this crate.
16//!
17//! This was probably a historical mistake in Rust library,
18//! that `retain` should do this at the very beginning.
19//! See [rust-lang/rust#25477](https://github.com/rust-lang/rust/issues/25477).
20//!
21//! From Rust 1.58, an unstable `retain_mut` method has been added to the std, see
22//! [rust-lang/rust#90829](https://github.com/rust-lang/rust/issues/90829).
23//! Once it gets stabilized, you can simply remove this crate.
24//!
25//! ## Examples
26//!
27//! ### `Vec`
28//!
29//! ```
30//! # use retain_mut::RetainMut;
31//! let mut vec = vec![1, 2, 3, 4];
32//! vec.retain_mut(|x| { *x *= 3; *x % 2 == 0 });
33//! assert_eq!(vec, [6, 12]);
34//! ```
35//!
36//! ### `VecDeque`
37//!
38//! ```
39//! # use retain_mut::RetainMut;
40//! # use std::collections::VecDeque;
41//! let mut deque = VecDeque::from(vec![1, 2, 3, 4]);
42//! deque.retain_mut(|x| { *x *= 3; *x % 2 == 0 });
43//! assert_eq!(deque, [6, 12]);
44//! ```
45
46#![no_std]
47
48extern crate alloc;
49
50use alloc::collections::vec_deque::VecDeque;
51use alloc::vec::Vec;
52use core::ptr;
53
54/// Trait that provides `retain_mut` method.
55#[deprecated = "Rust 1.61 has included retain_mut directly"]
56pub trait RetainMut<T> {
57    /// Retains only the elements specified by the predicate, passing a mutable reference to it.
58    ///
59    /// In other words, remove all elements `e` such that `f(&mut e)` returns `false`.
60    /// This method operates in place, visiting each element exactly once in the
61    /// original order, and preserves the order of the retained elements.
62    fn retain_mut<F>(&mut self, f: F)
63    where
64        F: FnMut(&mut T) -> bool;
65}
66
67impl<T> RetainMut<T> for Vec<T> {
68    // The implementation is based on
69    // https://github.com/rust-lang/rust/blob/03c8ffaacb040a8753ef8e1accea701bc9f5be85/library/alloc/src/vec/mod.rs#L1478-L1569
70    fn retain_mut<F>(&mut self, mut f: F)
71    where
72        F: FnMut(&mut T) -> bool,
73    {
74        let original_len = self.len();
75        // Avoid double drop if the drop guard is not executed,
76        // since we may make some holes during the process.
77        unsafe { self.set_len(0) };
78
79        // Vec: [Kept, Kept, Hole, Hole, Hole, Hole, Unchecked, Unchecked]
80        //      |<-              processed len   ->| ^- next to check
81        //                  |<-  deleted cnt     ->|
82        //      |<-              original_len                          ->|
83        // Kept: Elements which predicate returns true on.
84        // Hole: Moved or dropped element slot.
85        // Unchecked: Unchecked valid elements.
86        //
87        // This drop guard will be invoked when predicate or `drop` of element panicked.
88        // It shifts unchecked elements to cover holes and `set_len` to the correct length.
89        // In cases when predicate and `drop` never panick, it will be optimized out.
90        struct BackshiftOnDrop<'a, T> {
91            v: &'a mut Vec<T>,
92            processed_len: usize,
93            deleted_cnt: usize,
94            original_len: usize,
95        }
96
97        impl<T> Drop for BackshiftOnDrop<'_, T> {
98            fn drop(&mut self) {
99                if self.deleted_cnt > 0 {
100                    // SAFETY: Trailing unchecked items must be valid since we never touch them.
101                    unsafe {
102                        ptr::copy(
103                            self.v.as_ptr().add(self.processed_len),
104                            self.v
105                                .as_mut_ptr()
106                                .add(self.processed_len - self.deleted_cnt),
107                            self.original_len - self.processed_len,
108                        );
109                    }
110                }
111                // SAFETY: After filling holes, all items are in contiguous memory.
112                unsafe {
113                    self.v.set_len(self.original_len - self.deleted_cnt);
114                }
115            }
116        }
117
118        let mut g = BackshiftOnDrop {
119            v: self,
120            processed_len: 0,
121            deleted_cnt: 0,
122            original_len,
123        };
124
125        fn process_loop<F, T, const DELETED: bool>(
126            original_len: usize,
127            f: &mut F,
128            g: &mut BackshiftOnDrop<'_, T>,
129        ) where
130            F: FnMut(&mut T) -> bool,
131        {
132            while g.processed_len != original_len {
133                // SAFETY: Unchecked element must be valid.
134                let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) };
135                if !f(cur) {
136                    // Advance early to avoid double drop if `drop_in_place` panicked.
137                    g.processed_len += 1;
138                    g.deleted_cnt += 1;
139                    // SAFETY: We never touch this element again after dropped.
140                    unsafe { ptr::drop_in_place(cur) };
141                    // We already advanced the counter.
142                    if DELETED {
143                        continue;
144                    } else {
145                        break;
146                    }
147                }
148                if DELETED {
149                    // SAFETY: `deleted_cnt` > 0, so the hole slot must not overlap with current element.
150                    // We use copy for move, and never touch this element again.
151                    unsafe {
152                        let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt);
153                        ptr::copy_nonoverlapping(cur, hole_slot, 1);
154                    }
155                }
156                g.processed_len += 1;
157            }
158        }
159
160        // Stage 1: Nothing was deleted.
161        process_loop::<F, T, false>(original_len, &mut f, &mut g);
162
163        // Stage 2: Some elements were deleted.
164        process_loop::<F, T, true>(original_len, &mut f, &mut g);
165
166        // All item are processed. This can be optimized to `set_len` by LLVM.
167        drop(g);
168    }
169}
170
171impl<T> RetainMut<T> for VecDeque<T> {
172    // The implementation is based on
173    // https://github.com/rust-lang/rust/blob/3e21768a0a3fc84befd1cbe825ae6849e9941b73/library/alloc/src/collections/vec_deque/mod.rs#L2148-L2180
174    fn retain_mut<F>(&mut self, mut f: F)
175    where
176        F: FnMut(&mut T) -> bool,
177    {
178        let len = self.len();
179        let mut idx = 0;
180        let mut cur = 0;
181
182        // Stage 1: All values are retained.
183        while cur < len {
184            if !f(&mut self[cur]) {
185                cur += 1;
186                break;
187            }
188            cur += 1;
189            idx += 1;
190        }
191        // Stage 2: Swap retained value into current idx.
192        while cur < len {
193            if !f(&mut self[cur]) {
194                cur += 1;
195                continue;
196            }
197
198            self.swap(idx, cur);
199            cur += 1;
200            idx += 1;
201        }
202        // Stage 3: Trancate all values after idx.
203        if cur != idx {
204            self.truncate(idx);
205        }
206    }
207}