Optimize unnecessary check in VecDeque::retain

Signed-off-by: Xuanwo <github@xuanwo.io>

diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs
index 9a2205420a14b095781c2d06fff6300b91cbeb98..e4b28204158d9d3372ade952b5487eb1e96ebfd7 100644 (file)
--- a/library/alloc/src/collections/vec_deque/mod.rs
+++ b/library/alloc/src/collections/vec_deque/mod.rs
@@ -2129,16 +2129,32 @@ pub fn retain<F>(&mut self, mut f: F)
         F: FnMut(&T) -> bool,
     {
         let len = self.len();
-        let mut del = 0;
-        for i in 0..len {
-            if !f(&self[i]) {
-                del += 1;
-            } else if del > 0 {
-                self.swap(i - del, i);
+        let mut idx = 0;
+        let mut cur = 0;
+
+        // Stage 1: All values are retained.
+        while cur < len {
+            if !f(&self[cur]) {
+                cur += 1;
+                break;
             }
+            cur += 1;
+            idx += 1;
         }
-        if del > 0 {
-            self.truncate(len - del);
+        // Stage 2: Swap retained value into current idx.
+        while cur < len {
+            if !f(&self[cur]) {
+                cur += 1;
+                continue;
+            }
+
+            self.swap(idx, cur);
+            cur += 1;
+            idx += 1;
+        }
+        // Stage 3: Trancate all values after idx.
+        if cur != idx {
+            self.truncate(idx);
         }
     }
 

diff --git a/library/alloc/src/collections/vec_deque/tests.rs b/library/alloc/src/collections/vec_deque/tests.rs
index 6116cfe1d0110bdd33e0087e712a100ce5a066e7..bf1c73eb0780340261a37243b9320bb755fdd466 100644 (file)
--- a/library/alloc/src/collections/vec_deque/tests.rs
+++ b/library/alloc/src/collections/vec_deque/tests.rs
@@ -40,6 +40,39 @@ fn bench_pop_back_100(b: &mut test::Bencher) {
     })
 }
 
+#[bench]
+#[cfg_attr(miri, ignore)] // isolated Miri does not support benchmarks
+fn bench_retain_whole_10000(b: &mut test::Bencher) {
+    let v = (1..100000).collect::<VecDeque<u32>>();
+
+    b.iter(|| {
+        let mut v = v.clone();
+        v.retain(|x| *x > 0)
+    })
+}
+
+#[bench]
+#[cfg_attr(miri, ignore)] // isolated Miri does not support benchmarks
+fn bench_retain_odd_10000(b: &mut test::Bencher) {
+    let v = (1..100000).collect::<VecDeque<u32>>();
+
+    b.iter(|| {
+        let mut v = v.clone();
+        v.retain(|x| x & 1 == 0)
+    })
+}
+
+#[bench]
+#[cfg_attr(miri, ignore)] // isolated Miri does not support benchmarks
+fn bench_retain_half_10000(b: &mut test::Bencher) {
+    let v = (1..100000).collect::<VecDeque<u32>>();
+
+    b.iter(|| {
+        let mut v = v.clone();
+        v.retain(|x| *x > 50000)
+    })
+}
+
 #[bench]
 #[cfg_attr(miri, ignore)] // isolated Miri does not support benchmarks
 fn bench_pop_front_100(b: &mut test::Bencher) {
@@ -54,6 +87,8 @@ fn bench_pop_front_100(b: &mut test::Bencher) {
     })
 }
 
+
+
 #[test]
 fn test_swap_front_back_remove() {
     fn test(back: bool) {