rust/alloc/vec/drain_filter.rs

   1 // SPDX-License-Identifier: Apache-2.0 OR MIT
   2
   3 use crate::alloc::{Allocator, Global};
   4 use core::mem::{ManuallyDrop, SizedTypeProperties};
   5 use core::ptr;
   6 use core::slice;
   7
   8 use super::Vec;
   9
  10 /// An iterator which uses a closure to determine if an element should be removed.
  11 ///
  12 /// This struct is created by [`Vec::drain_filter`].
  13 /// See its documentation for more.
  14 ///
  15 /// # Example
  16 ///
  17 /// ```
  18 /// #![feature(drain_filter)]
  19 ///
  20 /// let mut v = vec![0, 1, 2];
  21 /// let iter: std::vec::DrainFilter<'_, _, _> = v.drain_filter(|x| *x % 2 == 0);
  22 /// ```
  23 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
  24 #[derive(Debug)]
  25 pub struct DrainFilter<
  26     'a,
  27     T,
  28     F,
  29     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
  30 > where
  31     F: FnMut(&mut T) -> bool,
  32 {
  33     pub(super) vec: &'a mut Vec<T, A>,
  34     /// The index of the item that will be inspected by the next call to `next`.
  35     pub(super) idx: usize,
  36     /// The number of items that have been drained (removed) thus far.
  37     pub(super) del: usize,
  38     /// The original length of `vec` prior to draining.
  39     pub(super) old_len: usize,
  40     /// The filter test predicate.
  41     pub(super) pred: F,
  42     /// A flag that indicates a panic has occurred in the filter test predicate.
  43     /// This is used as a hint in the drop implementation to prevent consumption
  44     /// of the remainder of the `DrainFilter`. Any unprocessed items will be
  45     /// backshifted in the `vec`, but no further items will be dropped or
  46     /// tested by the filter predicate.
  47     pub(super) panic_flag: bool,
  48 }
  49
  50 impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
  51 where
  52     F: FnMut(&mut T) -> bool,
  53 {
  54     /// Returns a reference to the underlying allocator.
  55     #[unstable(feature = "allocator_api", issue = "32838")]
  56     #[inline]
  57     pub fn allocator(&self) -> &A {
  58         self.vec.allocator()
  59     }
  60
  61     /// Keep unyielded elements in the source `Vec`.
  62     ///
  63     /// # Examples
  64     ///
  65     /// ```
  66     /// #![feature(drain_filter)]
  67     /// #![feature(drain_keep_rest)]
  68     ///
  69     /// let mut vec = vec!['a', 'b', 'c'];
  70     /// let mut drain = vec.drain_filter(|_| true);
  71     ///
  72     /// assert_eq!(drain.next().unwrap(), 'a');
  73     ///
  74     /// // This call keeps 'b' and 'c' in the vec.
  75     /// drain.keep_rest();
  76     ///
  77     /// // If we wouldn't call `keep_rest()`,
  78     /// // `vec` would be empty.
  79     /// assert_eq!(vec, ['b', 'c']);
  80     /// ```
  81     #[unstable(feature = "drain_keep_rest", issue = "101122")]
  82     pub fn keep_rest(self) {
  83         // At this moment layout looks like this:
  84         //
  85         //  _____________________/-- old_len
  86         // /                     \
  87         // [kept] [yielded] [tail]
  88         //        \_______/ ^-- idx
  89         //                \-- del
  90         //
  91         // Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
  92         //
  93         // 1. Move [tail] after [kept]
  94         // 2. Update length of the original vec to `old_len - del`
  95         //    a. In case of ZST, this is the only thing we want to do
  96         // 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
  97         let mut this = ManuallyDrop::new(self);
  98
  99         unsafe {
 100             // ZSTs have no identity, so we don't need to move them around.
 101             if !T::IS_ZST && this.idx < this.old_len && this.del > 0 {
 102                 let ptr = this.vec.as_mut_ptr();
 103                 let src = ptr.add(this.idx);
 104                 let dst = src.sub(this.del);
 105                 let tail_len = this.old_len - this.idx;
 106                 src.copy_to(dst, tail_len);
 107             }
 108
 109             let new_len = this.old_len - this.del;
 110             this.vec.set_len(new_len);
 111         }
 112     }
 113 }
 114
 115 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
 116 impl<T, F, A: Allocator> Iterator for DrainFilter<'_, T, F, A>
 117 where
 118     F: FnMut(&mut T) -> bool,
 119 {
 120     type Item = T;
 121
 122     fn next(&mut self) -> Option<T> {
 123         unsafe {
 124             while self.idx < self.old_len {
 125                 let i = self.idx;
 126                 let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
 127                 self.panic_flag = true;
 128                 let drained = (self.pred)(&mut v[i]);
 129                 self.panic_flag = false;
 130                 // Update the index *after* the predicate is called. If the index
 131                 // is updated prior and the predicate panics, the element at this
 132                 // index would be leaked.
 133                 self.idx += 1;
 134                 if drained {
 135                     self.del += 1;
 136                     return Some(ptr::read(&v[i]));
 137                 } else if self.del > 0 {
 138                     let del = self.del;
 139                     let src: *const T = &v[i];
 140                     let dst: *mut T = &mut v[i - del];
 141                     ptr::copy_nonoverlapping(src, dst, 1);
 142                 }
 143             }
 144             None
 145         }
 146     }
 147
 148     fn size_hint(&self) -> (usize, Option<usize>) {
 149         (0, Some(self.old_len - self.idx))
 150     }
 151 }
 152
 153 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
 154 impl<T, F, A: Allocator> Drop for DrainFilter<'_, T, F, A>
 155 where
 156     F: FnMut(&mut T) -> bool,
 157 {
 158     fn drop(&mut self) {
 159         struct BackshiftOnDrop<'a, 'b, T, F, A: Allocator>
 160         where
 161             F: FnMut(&mut T) -> bool,
 162         {
 163             drain: &'b mut DrainFilter<'a, T, F, A>,
 164         }
 165
 166         impl<'a, 'b, T, F, A: Allocator> Drop for BackshiftOnDrop<'a, 'b, T, F, A>
 167         where
 168             F: FnMut(&mut T) -> bool,
 169         {
 170             fn drop(&mut self) {
 171                 unsafe {
 172                     if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
 173                         // This is a pretty messed up state, and there isn't really an
 174                         // obviously right thing to do. We don't want to keep trying
 175                         // to execute `pred`, so we just backshift all the unprocessed
 176                         // elements and tell the vec that they still exist. The backshift
 177                         // is required to prevent a double-drop of the last successfully
 178                         // drained item prior to a panic in the predicate.
 179                         let ptr = self.drain.vec.as_mut_ptr();
 180                         let src = ptr.add(self.drain.idx);
 181                         let dst = src.sub(self.drain.del);
 182                         let tail_len = self.drain.old_len - self.drain.idx;
 183                         src.copy_to(dst, tail_len);
 184                     }
 185                     self.drain.vec.set_len(self.drain.old_len - self.drain.del);
 186                 }
 187             }
 188         }
 189
 190         let backshift = BackshiftOnDrop { drain: self };
 191
 192         // Attempt to consume any remaining elements if the filter predicate
 193         // has not yet panicked. We'll backshift any remaining elements
 194         // whether we've already panicked or if the consumption here panics.
 195         if !backshift.drain.panic_flag {
 196             backshift.drain.for_each(drop);
 197         }
 198     }
 199 }