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     #[unstable(feature = "drain_keep_rest", issue = "none")]
  61     pub fn keep_rest(self) {
  62         // At this moment layout looks like this:
  63         //
  64         //  _____________________/-- old_len
  65         // /                     \
  66         // [kept] [yielded] [tail]
  67         //        \_______/ ^-- idx
  68         //                \-- del
  69         //
  70         // Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
  71         //
  72         // 1. Move [tail] after [kept]
  73         // 2. Update length of the original vec to `old_len - del`
  74         //    a. In case of ZST, this is the only thing we want to do
  75         // 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
  76         let mut this = ManuallyDrop::new(self);
  77
  78         unsafe {
  79             // ZSTs have no identity, so we don't need to move them around.
  80             let needs_move = mem::size_of::<T>() != 0;
  81
  82             if needs_move && this.idx < this.old_len && this.del > 0 {
  83                 let ptr = this.vec.as_mut_ptr();
  84                 let src = ptr.add(this.idx);
  85                 let dst = src.sub(this.del);
  86                 let tail_len = this.old_len - this.idx;
  87                 src.copy_to(dst, tail_len);
  88             }
  89
  90             let new_len = this.old_len - this.del;
  91             this.vec.set_len(new_len);
  92         }
  93     }
  94 }
  95
  96 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
  97 impl<T, F, A: Allocator> Iterator for DrainFilter<'_, T, F, A>
  98 where
  99     F: FnMut(&mut T) -> bool,
 100 {
 101     type Item = T;
 102
 103     fn next(&mut self) -> Option<T> {
 104         unsafe {
 105             while self.idx < self.old_len {
 106                 let i = self.idx;
 107                 let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
 108                 self.panic_flag = true;
 109                 let drained = (self.pred)(&mut v[i]);
 110                 self.panic_flag = false;
 111                 // Update the index *after* the predicate is called. If the index
 112                 // is updated prior and the predicate panics, the element at this
 113                 // index would be leaked.
 114                 self.idx += 1;
 115                 if drained {
 116                     self.del += 1;
 117                     return Some(ptr::read(&v[i]));
 118                 } else if self.del > 0 {
 119                     let del = self.del;
 120                     let src: *const T = &v[i];
 121                     let dst: *mut T = &mut v[i - del];
 122                     ptr::copy_nonoverlapping(src, dst, 1);
 123                 }
 124             }
 125             None
 126         }
 127     }
 128
 129     fn size_hint(&self) -> (usize, Option<usize>) {
 130         (0, Some(self.old_len - self.idx))
 131     }
 132 }
 133
 134 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
 135 impl<T, F, A: Allocator> Drop for DrainFilter<'_, T, F, A>
 136 where
 137     F: FnMut(&mut T) -> bool,
 138 {
 139     fn drop(&mut self) {
 140         struct BackshiftOnDrop<'a, 'b, T, F, A: Allocator>
 141         where
 142             F: FnMut(&mut T) -> bool,
 143         {
 144             drain: &'b mut DrainFilter<'a, T, F, A>,
 145         }
 146
 147         impl<'a, 'b, T, F, A: Allocator> Drop for BackshiftOnDrop<'a, 'b, T, F, A>
 148         where
 149             F: FnMut(&mut T) -> bool,
 150         {
 151             fn drop(&mut self) {
 152                 unsafe {
 153                     if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
 154                         // This is a pretty messed up state, and there isn't really an
 155                         // obviously right thing to do. We don't want to keep trying
 156                         // to execute `pred`, so we just backshift all the unprocessed
 157                         // elements and tell the vec that they still exist. The backshift
 158                         // is required to prevent a double-drop of the last successfully
 159                         // drained item prior to a panic in the predicate.
 160                         let ptr = self.drain.vec.as_mut_ptr();
 161                         let src = ptr.add(self.drain.idx);
 162                         let dst = src.sub(self.drain.del);
 163                         let tail_len = self.drain.old_len - self.drain.idx;
 164                         src.copy_to(dst, tail_len);
 165                     }
 166                     self.drain.vec.set_len(self.drain.old_len - self.drain.del);
 167                 }
 168             }
 169         }
 170
 171         let backshift = BackshiftOnDrop { drain: self };
 172
 173         // Attempt to consume any remaining elements if the filter predicate
 174         // has not yet panicked. We'll backshift any remaining elements
 175         // whether we've already panicked or if the consumption here panics.
 176         if !backshift.drain.panic_flag {
 177             backshift.drain.for_each(drop);
 178         }
 179     }
 180 }