Rollup merge of rust-lang#147780 - tisonkun:vec-deque-extract-if, r=joboet

Implement VecDeque::extract_if This refers to rust-lang#147750.

Author: matthiaskrgr

library/alloc/src/collections/vec_deque/extract_if.rs

@@ -0,0 +1,149 @@
+use core::ops::{Range, RangeBounds};
+use core::{fmt, ptr, slice};
+
+use super::VecDeque;
+use crate::alloc::{Allocator, Global};
+
+/// An iterator which uses a closure to determine if an element should be removed.
+///
+/// This struct is created by [`VecDeque::extract_if`].
+/// See its documentation for more.
+///
+/// # Example
+///
+/// ```
+/// #![feature(vec_deque_extract_if)]
+///
+/// use std::collections::vec_deque::ExtractIf;
+/// use std::collections::vec_deque::VecDeque;
+///
+/// let mut v = VecDeque::from([0, 1, 2]);
+/// let iter: ExtractIf<'_, _, _> = v.extract_if(.., |x| *x % 2 == 0);
+/// ```
+#[unstable(feature = "vec_deque_extract_if", issue = "147750")]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+pub struct ExtractIf<
+ 'a,
+ T,
+ F,
+ #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
+> {
+ vec: &'a mut VecDeque<T, A>,
+ /// The index of the item that will be inspected by the next call to `next`.
+ idx: usize,
+ /// Elements at and beyond this point will be retained. Must be equal or smaller than `old_len`.
+ end: usize,
+ /// The number of items that have been drained (removed) thus far.
+ del: usize,
+ /// The original length of `vec` prior to draining.
+ old_len: usize,
+ /// The filter test predicate.
+ pred: F,
+}
+
+impl<'a, T, F, A: Allocator> ExtractIf<'a, T, F, A> {
+ pub(super) fn new<R: RangeBounds<usize>>(
+ vec: &'a mut VecDeque<T, A>,
+ pred: F,
+ range: R,
+ ) -> Self {
+ let old_len = vec.len();
+ let Range { start, end } = slice::range(range, ..old_len);
+
+ // Guard against the deque getting leaked (leak amplification)
+ vec.len = 0;
+ ExtractIf { vec, idx: start, del: 0, end, old_len, pred }
+ }
+
+ /// Returns a reference to the underlying allocator.
+ #[unstable(feature = "allocator_api", issue = "32838")]
+ #[inline]
+ pub fn allocator(&self) -> &A {
+ self.vec.allocator()
+ }
+}
+
+#[unstable(feature = "vec_deque_extract_if", issue = "147750")]
+impl<T, F, A: Allocator> Iterator for ExtractIf<'_, T, F, A>
+where
+ F: FnMut(&mut T) -> bool,
+{
+ type Item = T;
+
+ fn next(&mut self) -> Option<T> {
+ while self.idx < self.end {
+ let i = self.idx;
+ // SAFETY:
+ // We know that `i < self.end` from the if guard and that `self.end <= self.old_len` from
+ // the validity of `Self`. Therefore `i` points to an element within `vec`.
+ //
+ // Additionally, the i-th element is valid because each element is visited at most once
+ // and it is the first time we access vec[i].
+ //
+ // Note: we can't use `vec.get_mut(i).unwrap()` here since the precondition for that
+ // function is that i < vec.len, but we've set vec's length to zero.
+ let idx = self.vec.to_physical_idx(i);
+ let cur = unsafe { &mut *self.vec.ptr().add(idx) };
+ let drained = (self.pred)(cur);
+ // Update the index *after* the predicate is called. If the index
+ // is updated prior and the predicate panics, the element at this
+ // index would be leaked.
+ self.idx += 1;
+ if drained {
+ self.del += 1;
+ // SAFETY: We never touch this element again after returning it.
+ return Some(unsafe { ptr::read(cur) });
+ } else if self.del > 0 {
+ let hole_slot = self.vec.to_physical_idx(i - self.del);
+ // SAFETY: `self.del` > 0, so the hole slot must not overlap with current element.
+ // We use copy for move, and never touch this element again.
+ unsafe { self.vec.wrap_copy(idx, hole_slot, 1) };
+ }
+ }
+ None
+ }
+
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ (0, Some(self.end - self.idx))
+ }
+}
+
+#[unstable(feature = "vec_deque_extract_if", issue = "147750")]
+impl<T, F, A: Allocator> Drop for ExtractIf<'_, T, F, A> {
+ fn drop(&mut self) {
+ if self.del > 0 {
+ let src = self.vec.to_physical_idx(self.idx);
+ let dst = self.vec.to_physical_idx(self.idx - self.del);
+ let len = self.old_len - self.idx;
+ // SAFETY: Trailing unchecked items must be valid since we never touch them.
+ unsafe { self.vec.wrap_copy(src, dst, len) };
+ }
+ self.vec.len = self.old_len - self.del;
+ }
+}
+
+#[unstable(feature = "vec_deque_extract_if", issue = "147750")]
+impl<T, F, A> fmt::Debug for ExtractIf<'_, T, F, A>
+where
+ T: fmt::Debug,
+ A: Allocator,
+{
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ let peek = if self.idx < self.end {
+ let idx = self.vec.to_physical_idx(self.idx);
+ // This has to use pointer arithmetic as `self.vec[self.idx]` or
+ // `self.vec.get_unchecked(self.idx)` wouldn't work since we
+ // temporarily set the length of `self.vec` to zero.
+ //
+ // SAFETY:
+ // Since `self.idx` is smaller than `self.end` and `self.end` is
+ // smaller than `self.old_len`, `idx` is valid for indexing the
+ // buffer. Also, per the invariant of `self.idx`, this element
+ // has not been inspected/moved out yet.
+ Some(unsafe { &*self.vec.ptr().add(idx) })
+ } else {
+ None
+ };
+ f.debug_struct("ExtractIf").field("peek", &peek).finish_non_exhaustive()
+ }
+}

library/alloc/src/collections/vec_deque/mod.rs

@@ -32,6 +32,11 @@ pub use self::drain::Drain;
 
 mod drain;
 
+#[unstable(feature = "vec_deque_extract_if", issue = "147750")]
+pub use self::extract_if::ExtractIf;
+
+mod extract_if;
+
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::iter_mut::IterMut;
 
@@ -614,6 +619,95 @@ impl<T, A: Allocator> VecDeque<T, A> {
  }
  debug_assert!(self.head < self.capacity() || self.capacity() == 0);
  }
+
+ /// Creates an iterator which uses a closure to determine if an element in the range should be removed.
+ ///
+ /// If the closure returns `true`, the element is removed from the deque and yielded. If the closure
+ /// returns `false`, or panics, the element remains in the deque and will not be yielded.
+ ///
+ /// Only elements that fall in the provided range are considered for extraction, but any elements
+ /// after the range will still have to be moved if any element has been extracted.
+ ///
+ /// If the returned `ExtractIf` is not exhausted, e.g. because it is dropped without iterating
+ /// or the iteration short-circuits, then the remaining elements will be retained.
+ /// Use [`retain_mut`] with a negated predicate if you do not need the returned iterator.
+ ///
+ /// [`retain_mut`]: VecDeque::retain_mut
+ ///
+ /// Using this method is equivalent to the following code:
+ ///
+ /// ```
+ /// #![feature(vec_deque_extract_if)]
+ /// # use std::collections::VecDeque;
+ /// # let some_predicate = |x: &mut i32| { *x % 2 == 1 };
+ /// # let mut deq: VecDeque<_> = (0..10).collect();
+ /// # let mut deq2 = deq.clone();
+ /// # let range = 1..5;
+ /// let mut i = range.start;
+ /// let end_items = deq.len() - range.end;
+ /// # let mut extracted = vec![];
+ ///
+ /// while i < deq.len() - end_items {
+ /// if some_predicate(&mut deq[i]) {
+ /// let val = deq.remove(i).unwrap();
+ /// // your code here
+ /// # extracted.push(val);
+ /// } else {
+ /// i += 1;
+ /// }
+ /// }
+ ///
+ /// # let extracted2: Vec<_> = deq2.extract_if(range, some_predicate).collect();
+ /// # assert_eq!(deq, deq2);
+ /// # assert_eq!(extracted, extracted2);
+ /// ```
+ ///
+ /// But `extract_if` is easier to use. `extract_if` is also more efficient,
+ /// because it can backshift the elements of the array in bulk.
+ ///
+ /// The iterator also lets you mutate the value of each element in the
+ /// closure, regardless of whether you choose to keep or remove it.
+ ///
+ /// # Panics
+ ///
+ /// If `range` is out of bounds.
+ ///
+ /// # Examples
+ ///
+ /// Splitting a deque into even and odd values, reusing the original deque:
+ ///
+ /// ```
+ /// #![feature(vec_deque_extract_if)]
+ /// use std::collections::VecDeque;
+ ///
+ /// let mut numbers = VecDeque::from([1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
+ ///
+ /// let evens = numbers.extract_if(.., |x| *x % 2 == 0).collect::<VecDeque<_>>();
+ /// let odds = numbers;
+ ///
+ /// assert_eq!(evens, VecDeque::from([2, 4, 6, 8, 14]));
+ /// assert_eq!(odds, VecDeque::from([1, 3, 5, 9, 11, 13, 15]));
+ /// ```
+ ///
+ /// Using the range argument to only process a part of the deque:
+ ///
+ /// ```
+ /// #![feature(vec_deque_extract_if)]
+ /// use std::collections::VecDeque;
+ ///
+ /// let mut items = VecDeque::from([0, 0, 0, 0, 0, 0, 0, 1, 2, 1, 2, 1, 2]);
+ /// let ones = items.extract_if(7.., |x| *x == 1).collect::<VecDeque<_>>();
+ /// assert_eq!(items, VecDeque::from([0, 0, 0, 0, 0, 0, 0, 2, 2, 2]));
+ /// assert_eq!(ones.len(), 3);
+ /// ```
+ #[unstable(feature = "vec_deque_extract_if", issue = "147750")]
+ pub fn extract_if<F, R>(&mut self, range: R, filter: F) -> ExtractIf<'_, T, F, A>
+ where
+ F: FnMut(&mut T) -> bool,
+ R: RangeBounds<usize>,
+ {
+ ExtractIf::new(self, filter, range)
+ }
 }
 
 impl<T> VecDeque<T> {