From bc21f19e70668dab3d2d1d5a8bdf35a2385079e0 Mon Sep 17 00:00:00 2001
From: Niklas Fiekas <niklas.fiekas@backscattering.de>
Date: Fri, 8 Oct 2021 10:38:29 +0200
Subject: [PATCH] Follow optimization of std::Vec::retain
 (rust-lang/rust#88060)

---
 src/arrayvec.rs | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/arrayvec.rs b/src/arrayvec.rs
index 6b4faf8..1fe6adf 100644
--- a/src/arrayvec.rs
+++ b/src/arrayvec.rs
@@ -493,21 +493,38 @@ impl<T, const CAP: usize> ArrayVec<T, CAP> {
 
         let mut g = BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len };
 
-        while g.processed_len < original_len {
+        #[inline(always)]
+        fn process_one<F: FnMut(&mut T) -> bool, T, const CAP: usize, const DELETED: bool>(
+            f: &mut F,
+            g: &mut BackshiftOnDrop<'_, T, CAP>
+        ) -> bool {
             let cur = unsafe { g.v.as_mut_ptr().add(g.processed_len) };
             if !f(unsafe { &mut *cur }) {
                 g.processed_len += 1;
                 g.deleted_cnt += 1;
                 unsafe { ptr::drop_in_place(cur) };
-                continue;
+                return false;
             }
-            if g.deleted_cnt > 0 {
+            if DELETED {
                 unsafe {
                     let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt);
                     ptr::copy_nonoverlapping(cur, hole_slot, 1);
                 }
             }
             g.processed_len += 1;
+            true
+        }
+
+        // Stage 1: Nothing was deleted.
+        while g.processed_len != original_len {
+            if !process_one::<F, T, CAP, false>(&mut f, &mut g) {
+                break;
+            }
+        }
+
+        // Stage 2: Some elements were deleted.
+        while g.processed_len != original_len {
+            process_one::<F, T, CAP, true>(&mut f, &mut g);
         }
 
         drop(g);