Add retain_mut (#198)

Fuuzetsu · web-flow · commit 5097217dccab · 2024-07-05T12:56:31.000-06:00
* Add retain_mut This adds a `retain_mut` method to `ArrayVec` and `TinyVec` that's just `retain` but with mutable references. This method is present on `std::vec::Vec` as of a while ago so it's nice to have in `tinyvec` too. Fixes #195 * `TinyVec::retain_mut` requires rustc_1_61 feature It calls `std::vec::Vec::retain_mut` which is only available since Rust 1.61. People that still want to use a very old compiler (CI seems to use 1.47) still get to compile the crate, they just don't get `TinyVec::retain_mut`. Note that `ArrayVec::retain_mut` is implemented inside the crate so it's still available, just the `TinyVec` format is not. * Add some tests for ArrayVec::retain_mut
diff --git a/Cargo.toml b/Cargo.toml
@@ -41,6 +41,10 @@ rustc_1_55 = []
 # add try_reserve functions to types that heap allocate.
 rustc_1_57 = ["rustc_1_55"]
 
+# features that require rustc 1.61
+# add retain_mut function to TinyVec
+rustc_1_61 = []
+
 # allow use of nightly feature `slice_partition_dedup`,
 # will become useless once that is stabilized:
 # https://github.com/rust-lang/rust/issues/54279
diff --git a/src/arrayvec.rs b/src/arrayvec.rs
@@ -816,6 +816,63 @@ impl<A: Array> ArrayVec<A> {
     }
   }
 
+  /// Retains only the elements specified by the predicate, passing a mutable
+  /// reference to it.
+  ///
+  /// In other words, remove all elements e such that f(&mut e) returns false.
+  /// This method operates in place, visiting each element exactly once in the
+  /// original order, and preserves the order of the retained elements.
+  ///
+  ///
+  /// ## Example
+  ///
+  /// ```rust
+  /// # use tinyvec::*;
+  ///
+  /// let mut av = array_vec!([i32; 10] => 1, 1, 2, 3, 3, 4);
+  /// av.retain_mut(|x| if *x % 2 == 0 { *x *= 2; true } else { false });
+  /// assert_eq!(&av[..], [4, 8]);
+  /// ```
+  #[inline]
+  pub fn retain_mut<F>(&mut self, mut acceptable: F)
+  where
+    F: FnMut(&mut A::Item) -> bool,
+  {
+    // Drop guard to contain exactly the remaining elements when the test
+    // panics.
+    struct JoinOnDrop<'vec, Item> {
+      items: &'vec mut [Item],
+      done_end: usize,
+      // Start of tail relative to `done_end`.
+      tail_start: usize,
+    }
+
+    impl<Item> Drop for JoinOnDrop<'_, Item> {
+      fn drop(&mut self) {
+        self.items[self.done_end..].rotate_left(self.tail_start);
+      }
+    }
+
+    let mut rest = JoinOnDrop {
+      items: &mut self.data.as_slice_mut()[..self.len as usize],
+      done_end: 0,
+      tail_start: 0,
+    };
+
+    let len = self.len as usize;
+    for idx in 0..len {
+      // Loop start invariant: idx = rest.done_end + rest.tail_start
+      if !acceptable(&mut rest.items[idx]) {
+        let _ = core::mem::take(&mut rest.items[idx]);
+        self.len -= 1;
+        rest.tail_start += 1;
+      } else {
+        rest.items.swap(rest.done_end, idx);
+        rest.done_end += 1;
+      }
+    }
+  }
+
   /// Forces the length of the vector to `new_len`.
   ///
   /// ## Panics
@@ -1893,3 +1950,45 @@ where
     Ok(new_arrayvec)
   }
 }
+
+#[cfg(test)]
+mod test {
+  use super::*;
+
+  #[test]
+  fn retain_mut_empty_vec() {
+    let mut av: ArrayVec<[i32; 4]> = ArrayVec::new();
+    av.retain_mut(|&mut x| x % 2 == 0);
+    assert_eq!(av.len(), 0);
+  }
+
+  #[test]
+  fn retain_mut_all_elements() {
+    let mut av: ArrayVec<[i32; 4]> = array_vec!([i32; 4] => 2, 4, 6, 8);
+    av.retain_mut(|&mut x| x % 2 == 0);
+    assert_eq!(av.len(), 4);
+    assert_eq!(av.as_slice(), &[2, 4, 6, 8]);
+  }
+
+  #[test]
+  fn retain_mut_some_elements() {
+    let mut av: ArrayVec<[i32; 4]> = array_vec!([i32; 4] => 1, 2, 3, 4);
+    av.retain_mut(|&mut x| x % 2 == 0);
+    assert_eq!(av.len(), 2);
+    assert_eq!(av.as_slice(), &[2, 4]);
+  }
+
+  #[test]
+  fn retain_mut_no_elements() {
+    let mut av: ArrayVec<[i32; 4]> = array_vec!([i32; 4] => 1, 3, 5, 7);
+    av.retain_mut(|&mut x| x % 2 == 0);
+    assert_eq!(av.len(), 0);
+  }
+
+  #[test]
+  fn retain_mut_zero_capacity() {
+    let mut av: ArrayVec<[i32; 0]> = ArrayVec::new();
+    av.retain_mut(|&mut x| x % 2 == 0);
+    assert_eq!(av.len(), 0);
+  }
+}
diff --git a/src/tinyvec.rs b/src/tinyvec.rs
@@ -647,6 +647,27 @@ impl<A: Array> TinyVec<A> {
     }
   }
 
+  /// Walk the vec and keep only the elements that pass the predicate given,
+  /// having the opportunity to modify the elements at the same time.
+  ///
+  /// ## Example
+  ///
+  /// ```rust
+  /// use tinyvec::*;
+  ///
+  /// let mut tv = tiny_vec!([i32; 10] => 1, 2, 3, 4);
+  /// tv.retain_mut(|x| if *x % 2 == 0 { *x *= 2; true } else { false });
+  /// assert_eq!(tv.as_slice(), &[4, 8][..]);
+  /// ```
+  #[inline]
+  #[cfg(feature = "rustc_1_61")]
+  pub fn retain_mut<F: FnMut(&mut A::Item) -> bool>(&mut self, acceptable: F) {
+    match self {
+      TinyVec::Inline(i) => i.retain_mut(acceptable),
+      TinyVec::Heap(h) => h.retain_mut(acceptable),
+    }
+  }
+
   /// Helper for getting the mut slice.
   #[inline(always)]
   #[must_use]
