Optimize Vec::retain

[rust.git] / library / alloc / src / vec / mod.rs

diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index 1ca194c336112e0084ee4b66afe4db7bfb65701d..2c510b8b2ae61de27fc422433b60a23574110466 100644 (file)
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -1371,21 +1371,78 @@ pub fn retain<F>(&mut self, mut f: F)
         F: FnMut(&T) -> bool,
     {
         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);
+        // Avoid double drop if the drop guard is not executed,
+        // since we may make some holes during the process.
+        unsafe { self.set_len(0) };
+
+        // Vec: [Kept, Kept, Hole, Hole, Hole, Hole, Unchecked, Unchecked]
+        //      |<-              processed len   ->| ^- next to check
+        //                  |<-  deleted cnt     ->|
+        //      |<-              original_len                          ->|
+        // Kept: Elements which predicate returns true on.
+        // Hole: Moved or dropped element slot.
+        // Unchecked: Unchecked valid elements.
+        //
+        // This drop guard will be invoked when predicate or `drop` of element panicked.
+        // It shifts unchecked elements to cover holes and `set_len` to the correct length.
+        // In cases when predicate and `drop` never panick, it will be optimized out.
+        struct BackshiftOnDrop<'a, T, A: Allocator> {
+            v: &'a mut Vec<T, A>,
+            processed_len: usize,
+            deleted_cnt: usize,
+            original_len: usize,
+        }
+
+        impl<T, A: Allocator> Drop for BackshiftOnDrop<'_, T, A> {
+            fn drop(&mut self) {
+                if self.deleted_cnt > 0 {
+                    // SAFETY: Fill the hole of dropped or moved
+                    unsafe {
+                        ptr::copy(
+                            self.v.as_ptr().offset(self.processed_len as isize),
+                            self.v
+                                .as_mut_ptr()
+                                .offset(self.processed_len as isize - self.deleted_cnt as isize),
+                            self.original_len - self.processed_len,
+                        );
+                        self.v.set_len(self.original_len - self.deleted_cnt);
+                    }
                 }
             }
         }
-        if del > 0 {
-            self.truncate(len - del);
+
+        let mut guard = BackshiftOnDrop {
+            v: self,
+            processed_len: 0,
+            deleted_cnt: 0,
+            original_len: len,
+        };
+
+        let mut del = 0usize;
+        for i in 0..len {
+            // SAFETY: Unchecked element must be valid.
+            let cur = unsafe { &mut *guard.v.as_mut_ptr().offset(i as isize) };
+            if !f(cur) {
+                del += 1;
+                // Advance early to avoid double drop if `drop_in_place` panicked.
+                guard.processed_len = i + 1;
+                guard.deleted_cnt = del;
+                // SAFETY: We never touch this element again after dropped.
+                unsafe { ptr::drop_in_place(cur) };
+            } else if del > 0 {
+                // SAFETY: `del` > 0 so the hole slot must not overlap with current element.
+                // We use copy for move, and never touch this element again.
+                unsafe {
+                    let hole_slot = guard.v.as_mut_ptr().offset(i as isize - del as isize);
+                    ptr::copy_nonoverlapping(cur, hole_slot, 1);
+                }
+                guard.processed_len = i + 1;
+            }
         }
+
+        // All holes are at the end now. Simply cut them out.
+        unsafe { guard.v.set_len(len - del) };
+        mem::forget(guard);
     }
 
     /// Removes all but the first of consecutive elements in the vector that resolve to the same
@@ -1953,27 +2010,6 @@ pub fn dedup(&mut self) {
     }
 }
 
-impl<T, A: Allocator> Vec<T, A> {
-    /// Removes the first instance of `item` from the vector if the item exists.
-    ///
-    /// This method will be removed soon.
-    #[unstable(feature = "vec_remove_item", reason = "recently added", issue = "40062")]
-    #[rustc_deprecated(
-        reason = "Removing the first item equal to a needle is already easily possible \
-            with iterators and the current Vec methods. Furthermore, having a method for \
-            one particular case of removal (linear search, only the first item, no swap remove) \
-            but not for others is inconsistent. This method will be removed soon.",
-        since = "1.46.0"
-    )]
-    pub fn remove_item<V>(&mut self, item: &V) -> Option<T>
-    where
-        T: PartialEq<V>,
-    {
-        let pos = self.iter().position(|x| *x == *item)?;
-        Some(self.remove(pos))
-    }
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // Internal methods and functions
 ////////////////////////////////////////////////////////////////////////////////