library/alloc/src/vec/drain_filter.rs

   1 use crate::alloc::{Allocator, Global};
   2 use core::mem::{self, ManuallyDrop};
   3 use core::ptr;
   4 use core::slice;
   5
   6 use super::Vec;
   7
   8 /// An iterator which uses a closure to determine if an element should be removed.
   9 ///
  10 /// This struct is created by [`Vec::drain_filter`].
  11 /// See its documentation for more.
  12 ///
  13 /// # Example
  14 ///
  15 /// ```
  16 /// #![feature(drain_filter)]
  17 ///
  18 /// let mut v = vec![0, 1, 2];
  19 /// let iter: std::vec::DrainFilter<_, _> = v.drain_filter(|x| *x % 2 == 0);
  20 /// ```
  21 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
  22 #[derive(Debug)]
  23 pub struct DrainFilter<
  24     'a,
  25     T,
  26     F,
  27     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
  28 > where
  29     F: FnMut(&mut T) -> bool,
  30 {
  31     pub(super) vec: &'a mut Vec<T, A>,
  32     /// The index of the item that will be inspected by the next call to `next`.
  33     pub(super) idx: usize,
  34     /// The number of items that have been drained (removed) thus far.
  35     pub(super) del: usize,
  36     /// The original length of `vec` prior to draining.
  37     pub(super) old_len: usize,
  38     /// The filter test predicate.
  39     pub(super) pred: F,
  40     /// A flag that indicates a panic has occurred in the filter test predicate.
  41     /// This is used as a hint in the drop implementation to prevent consumption
  42     /// of the remainder of the `DrainFilter`. Any unprocessed items will be
  43     /// backshifted in the `vec`, but no further items will be dropped or
  44     /// tested by the filter predicate.
  45     pub(super) panic_flag: bool,
  46 }
  47
  48 impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
  49 where
  50     F: FnMut(&mut T) -> bool,
  51 {
  52     /// Returns a reference to the underlying allocator.
  53     #[unstable(feature = "allocator_api", issue = "32838")]
  54     #[inline]
  55     pub fn allocator(&self) -> &A {
  56         self.vec.allocator()
  57     }
  58
  59     /// Keep unyielded elements in the source `Vec`.
  60     ///
  61     /// # Examples
  62     ///
  63     /// ```
  64     /// #![feature(drain_filter)]
  65     /// #![feature(drain_keep_rest)]
  66     ///
  67     /// let mut vec = vec!['a', 'b', 'c'];
  68     /// let mut drain = vec.drain_filter(|_| true);
  69     ///
  70     /// assert_eq!(drain.next().unwrap(), 'a');
  71     ///
  72     /// // This call keeps 'b' and 'c' in the vec.
  73     /// drain.keep_rest();
  74     ///
  75     /// // If we wouldn't call `keep_rest()`,
  76     /// // `vec` would be empty.
  77     /// assert_eq!(vec, ['b', 'c']);
  78     /// ```
  79     #[unstable(feature = "drain_keep_rest", issue = "101122")]
  80     pub fn keep_rest(self) {
  81         // At this moment layout looks like this:
  82         //
  83         //  _____________________/-- old_len
  84         // /                     \
  85         // [kept] [yielded] [tail]
  86         //        \_______/ ^-- idx
  87         //                \-- del
  88         //
  89         // Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
  90         //
  91         // 1. Move [tail] after [kept]
  92         // 2. Update length of the original vec to `old_len - del`
  93         //    a. In case of ZST, this is the only thing we want to do
  94         // 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
  95         let mut this = ManuallyDrop::new(self);
  96
  97         unsafe {
  98             // ZSTs have no identity, so we don't need to move them around.
  99             let needs_move = mem::size_of::<T>() != 0;
 100
 101             if needs_move && 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 }