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] allocator(&self) -> &A57 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")] keep_rest(self)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 next(&mut self) -> Option<T>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 size_hint(&self) -> (usize, Option<usize>)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 { drop(&mut self)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 } 200
