diff options
Diffstat (limited to 'library/alloc/src/vec')
| -rw-r--r-- | library/alloc/src/vec/extract_if.rs | 50 | ||||
| -rw-r--r-- | library/alloc/src/vec/mod.rs | 43 |
2 files changed, 53 insertions, 40 deletions
diff --git a/library/alloc/src/vec/extract_if.rs b/library/alloc/src/vec/extract_if.rs index 72d51e89044..4db13981596 100644 --- a/library/alloc/src/vec/extract_if.rs +++ b/library/alloc/src/vec/extract_if.rs @@ -1,3 +1,4 @@ +use core::ops::{Range, RangeBounds}; use core::{ptr, slice}; use super::Vec; @@ -14,7 +15,7 @@ use crate::alloc::{Allocator, Global}; /// #![feature(extract_if)] /// /// let mut v = vec![0, 1, 2]; -/// let iter: std::vec::ExtractIf<'_, _, _> = v.extract_if(|x| *x % 2 == 0); +/// let iter: std::vec::ExtractIf<'_, _, _> = v.extract_if(.., |x| *x % 2 == 0); /// ``` #[unstable(feature = "extract_if", reason = "recently added", issue = "43244")] #[derive(Debug)] @@ -24,24 +25,32 @@ pub struct ExtractIf< T, F, #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global, -> where - F: FnMut(&mut T) -> bool, -{ - pub(super) vec: &'a mut Vec<T, A>, +> { + vec: &'a mut Vec<T, A>, /// The index of the item that will be inspected by the next call to `next`. - pub(super) idx: usize, + 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. - pub(super) del: usize, + del: usize, /// The original length of `vec` prior to draining. - pub(super) old_len: usize, + old_len: usize, /// The filter test predicate. - pub(super) pred: F, + pred: F, } -impl<T, F, A: Allocator> ExtractIf<'_, T, F, A> -where - F: FnMut(&mut T) -> bool, -{ +impl<'a, T, F, A: Allocator> ExtractIf<'a, T, F, A> { + pub(super) fn new<R: RangeBounds<usize>>(vec: &'a mut Vec<T, A>, pred: F, range: R) -> Self { + let old_len = vec.len(); + let Range { start, end } = slice::range(range, ..old_len); + + // Guard against the vec getting leaked (leak amplification) + unsafe { + vec.set_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] @@ -59,7 +68,7 @@ where fn next(&mut self) -> Option<T> { unsafe { - while self.idx < self.old_len { + while self.idx < self.end { let i = self.idx; let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len); let drained = (self.pred)(&mut v[i]); @@ -82,24 +91,15 @@ where } fn size_hint(&self) -> (usize, Option<usize>) { - (0, Some(self.old_len - self.idx)) + (0, Some(self.end - self.idx)) } } #[unstable(feature = "extract_if", reason = "recently added", issue = "43244")] -impl<T, F, A: Allocator> Drop for ExtractIf<'_, T, F, A> -where - F: FnMut(&mut T) -> bool, -{ +impl<T, F, A: Allocator> Drop for ExtractIf<'_, T, F, A> { fn drop(&mut self) { unsafe { if self.idx < self.old_len && self.del > 0 { - // This is a pretty messed up state, and there isn't really an - // obviously right thing to do. We don't want to keep trying - // to execute `pred`, so we just backshift all the unprocessed - // elements and tell the vec that they still exist. The backshift - // is required to prevent a double-drop of the last successfully - // drained item prior to a panic in the predicate. let ptr = self.vec.as_mut_ptr(); let src = ptr.add(self.idx); let dst = src.sub(self.del); diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 7e7a8ff72c7..5f4b85b58a9 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -3615,12 +3615,15 @@ impl<T, A: Allocator> Vec<T, A> { Splice { drain: self.drain(range), replace_with: replace_with.into_iter() } } - /// Creates an iterator which uses a closure to determine if an element should be removed. + /// Creates an iterator which uses a closure to determine if element in the range should be removed. /// /// If the closure returns true, then the element is removed and yielded. /// If the closure returns false, the element will remain in the vector and will not be yielded /// by the iterator. /// + /// 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`] with a negated predicate if you do not need the returned iterator. @@ -3630,10 +3633,12 @@ impl<T, A: Allocator> Vec<T, A> { /// Using this method is equivalent to the following code: /// /// ``` + /// # use std::cmp::min; /// # let some_predicate = |x: &mut i32| { *x == 2 || *x == 3 || *x == 6 }; /// # let mut vec = vec![1, 2, 3, 4, 5, 6]; - /// let mut i = 0; - /// while i < vec.len() { + /// # let range = 1..4; + /// let mut i = range.start; + /// while i < min(vec.len(), range.end) { /// if some_predicate(&mut vec[i]) { /// let val = vec.remove(i); /// // your code here @@ -3648,8 +3653,12 @@ impl<T, A: Allocator> Vec<T, A> { /// But `extract_if` is easier to use. `extract_if` is also more efficient, /// because it can backshift the elements of the array in bulk. /// - /// Note that `extract_if` also lets you mutate every element in the filter closure, - /// regardless of whether you choose to keep or remove it. + /// Note that `extract_if` also lets you mutate the elements passed to the filter closure, + /// regardless of whether you choose to keep or remove them. + /// + /// # Panics + /// + /// If `range` is out of bounds. /// /// # Examples /// @@ -3659,25 +3668,29 @@ impl<T, A: Allocator> Vec<T, A> { /// #![feature(extract_if)] /// let mut numbers = vec![1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]; /// - /// let evens = numbers.extract_if(|x| *x % 2 == 0).collect::<Vec<_>>(); + /// let evens = numbers.extract_if(.., |x| *x % 2 == 0).collect::<Vec<_>>(); /// let odds = numbers; /// /// assert_eq!(evens, vec![2, 4, 6, 8, 14]); /// assert_eq!(odds, vec![1, 3, 5, 9, 11, 13, 15]); /// ``` + /// + /// Using the range argument to only process a part of the vector: + /// + /// ``` + /// #![feature(extract_if)] + /// let mut items = vec![0, 0, 0, 0, 0, 0, 0, 1, 2, 1, 2, 1, 2]; + /// let ones = items.extract_if(7.., |x| *x == 1).collect::<Vec<_>>(); + /// assert_eq!(items, vec![0, 0, 0, 0, 0, 0, 0, 2, 2, 2]); + /// assert_eq!(ones.len(), 3); + /// ``` #[unstable(feature = "extract_if", reason = "recently added", issue = "43244")] - pub fn extract_if<F>(&mut self, filter: F) -> ExtractIf<'_, T, F, A> + 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>, { - let old_len = self.len(); - - // Guard against us getting leaked (leak amplification) - unsafe { - self.set_len(0); - } - - ExtractIf { vec: self, idx: 0, del: 0, old_len, pred: filter } + ExtractIf::new(self, filter, range) } } |