| 1 | use crate::intrinsics::unlikely; |
| 2 | use crate::iter::adapters::SourceIter; |
| 3 | use crate::iter::adapters::zip::try_get_unchecked; |
| 4 | use crate::iter::{ |
| 5 | FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, TrustedRandomAccess, |
| 6 | TrustedRandomAccessNoCoerce, |
| 7 | }; |
| 8 | use crate::num::NonZero; |
| 9 | use crate::ops::{ControlFlow, Try}; |
| 10 | |
| 11 | /// An iterator that skips over `n` elements of `iter`. |
| 12 | /// |
| 13 | /// This `struct` is created by the [`skip`] method on [`Iterator`]. See its |
| 14 | /// documentation for more. |
| 15 | /// |
| 16 | /// [`skip`]: Iterator::skip |
| 17 | /// [`Iterator`]: trait.Iterator.html |
| 18 | #[derive (Clone, Debug)] |
| 19 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
| 20 | #[stable (feature = "rust1" , since = "1.0.0" )] |
| 21 | pub struct Skip<I> { |
| 22 | iter: I, |
| 23 | n: usize, |
| 24 | } |
| 25 | |
| 26 | impl<I> Skip<I> { |
| 27 | pub(in crate::iter) fn new(iter: I, n: usize) -> Skip<I> { |
| 28 | Skip { iter, n } |
| 29 | } |
| 30 | } |
| 31 | |
| 32 | #[stable (feature = "rust1" , since = "1.0.0" )] |
| 33 | impl<I> Iterator for Skip<I> |
| 34 | where |
| 35 | I: Iterator, |
| 36 | { |
| 37 | type Item = <I as Iterator>::Item; |
| 38 | |
| 39 | #[inline ] |
| 40 | fn next(&mut self) -> Option<I::Item> { |
| 41 | if unlikely(self.n > 0) { |
| 42 | self.iter.nth(crate::mem::take(&mut self.n)) |
| 43 | } else { |
| 44 | self.iter.next() |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | #[inline ] |
| 49 | fn nth(&mut self, n: usize) -> Option<I::Item> { |
| 50 | if self.n > 0 { |
| 51 | let skip: usize = crate::mem::take(&mut self.n); |
| 52 | // Checked add to handle overflow case. |
| 53 | let n = match skip.checked_add(n) { |
| 54 | Some(nth) => nth, |
| 55 | None => { |
| 56 | // In case of overflow, load skip value, before loading `n`. |
| 57 | // Because the amount of elements to iterate is beyond `usize::MAX`, this |
| 58 | // is split into two `nth` calls where the `skip` `nth` call is discarded. |
| 59 | self.iter.nth(skip - 1)?; |
| 60 | n |
| 61 | } |
| 62 | }; |
| 63 | // Load nth element including skip. |
| 64 | self.iter.nth(n) |
| 65 | } else { |
| 66 | self.iter.nth(n) |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | #[inline ] |
| 71 | fn count(mut self) -> usize { |
| 72 | if self.n > 0 { |
| 73 | // nth(n) skips n+1 |
| 74 | if self.iter.nth(self.n - 1).is_none() { |
| 75 | return 0; |
| 76 | } |
| 77 | } |
| 78 | self.iter.count() |
| 79 | } |
| 80 | |
| 81 | #[inline ] |
| 82 | fn last(mut self) -> Option<I::Item> { |
| 83 | if self.n > 0 { |
| 84 | // nth(n) skips n+1 |
| 85 | self.iter.nth(self.n - 1)?; |
| 86 | } |
| 87 | self.iter.last() |
| 88 | } |
| 89 | |
| 90 | #[inline ] |
| 91 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 92 | let (lower, upper) = self.iter.size_hint(); |
| 93 | |
| 94 | let lower = lower.saturating_sub(self.n); |
| 95 | let upper = match upper { |
| 96 | Some(x) => Some(x.saturating_sub(self.n)), |
| 97 | None => None, |
| 98 | }; |
| 99 | |
| 100 | (lower, upper) |
| 101 | } |
| 102 | |
| 103 | #[inline ] |
| 104 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
| 105 | where |
| 106 | Self: Sized, |
| 107 | Fold: FnMut(Acc, Self::Item) -> R, |
| 108 | R: Try<Output = Acc>, |
| 109 | { |
| 110 | let n = self.n; |
| 111 | self.n = 0; |
| 112 | if n > 0 { |
| 113 | // nth(n) skips n+1 |
| 114 | if self.iter.nth(n - 1).is_none() { |
| 115 | return try { init }; |
| 116 | } |
| 117 | } |
| 118 | self.iter.try_fold(init, fold) |
| 119 | } |
| 120 | |
| 121 | #[inline ] |
| 122 | fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc |
| 123 | where |
| 124 | Fold: FnMut(Acc, Self::Item) -> Acc, |
| 125 | { |
| 126 | if self.n > 0 { |
| 127 | // nth(n) skips n+1 |
| 128 | if self.iter.nth(self.n - 1).is_none() { |
| 129 | return init; |
| 130 | } |
| 131 | } |
| 132 | self.iter.fold(init, fold) |
| 133 | } |
| 134 | |
| 135 | #[inline ] |
| 136 | #[rustc_inherit_overflow_checks ] |
| 137 | fn advance_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> { |
| 138 | let skip_inner = self.n; |
| 139 | let skip_and_advance = skip_inner.saturating_add(n); |
| 140 | |
| 141 | let remainder = match self.iter.advance_by(skip_and_advance) { |
| 142 | Ok(()) => 0, |
| 143 | Err(n) => n.get(), |
| 144 | }; |
| 145 | let advanced_inner = skip_and_advance - remainder; |
| 146 | n -= advanced_inner.saturating_sub(skip_inner); |
| 147 | self.n = self.n.saturating_sub(advanced_inner); |
| 148 | |
| 149 | // skip_and_advance may have saturated |
| 150 | if unlikely(remainder == 0 && n > 0) { |
| 151 | n = match self.iter.advance_by(n) { |
| 152 | Ok(()) => 0, |
| 153 | Err(n) => n.get(), |
| 154 | } |
| 155 | } |
| 156 | |
| 157 | NonZero::new(n).map_or(Ok(()), Err) |
| 158 | } |
| 159 | |
| 160 | #[doc (hidden)] |
| 161 | unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item |
| 162 | where |
| 163 | Self: TrustedRandomAccessNoCoerce, |
| 164 | { |
| 165 | // SAFETY: the caller must uphold the contract for |
| 166 | // `Iterator::__iterator_get_unchecked`. |
| 167 | // |
| 168 | // Dropping the skipped prefix when index 0 is passed is safe |
| 169 | // since |
| 170 | // * the caller passing index 0 means that the inner iterator has more items than `self.n` |
| 171 | // * TRA contract requires that get_unchecked will only be called once |
| 172 | // (unless elements are copyable) |
| 173 | // * it does not conflict with in-place iteration since index 0 must be accessed |
| 174 | // before something is written into the storage used by the prefix |
| 175 | unsafe { |
| 176 | if Self::MAY_HAVE_SIDE_EFFECT && idx == 0 { |
| 177 | for skipped_idx in 0..self.n { |
| 178 | drop(try_get_unchecked(&mut self.iter, skipped_idx)); |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | try_get_unchecked(&mut self.iter, idx + self.n) |
| 183 | } |
| 184 | } |
| 185 | } |
| 186 | |
| 187 | #[stable (feature = "rust1" , since = "1.0.0" )] |
| 188 | impl<I> ExactSizeIterator for Skip<I> where I: ExactSizeIterator {} |
| 189 | |
| 190 | #[stable (feature = "double_ended_skip_iterator" , since = "1.9.0" )] |
| 191 | impl<I> DoubleEndedIterator for Skip<I> |
| 192 | where |
| 193 | I: DoubleEndedIterator + ExactSizeIterator, |
| 194 | { |
| 195 | fn next_back(&mut self) -> Option<Self::Item> { |
| 196 | if self.len() > 0 { self.iter.next_back() } else { None } |
| 197 | } |
| 198 | |
| 199 | #[inline ] |
| 200 | fn nth_back(&mut self, n: usize) -> Option<I::Item> { |
| 201 | let len = self.len(); |
| 202 | if n < len { |
| 203 | self.iter.nth_back(n) |
| 204 | } else { |
| 205 | if len > 0 { |
| 206 | // consume the original iterator |
| 207 | self.iter.nth_back(len - 1); |
| 208 | } |
| 209 | None |
| 210 | } |
| 211 | } |
| 212 | |
| 213 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
| 214 | where |
| 215 | Self: Sized, |
| 216 | Fold: FnMut(Acc, Self::Item) -> R, |
| 217 | R: Try<Output = Acc>, |
| 218 | { |
| 219 | fn check<T, Acc, R: Try<Output = Acc>>( |
| 220 | mut n: usize, |
| 221 | mut fold: impl FnMut(Acc, T) -> R, |
| 222 | ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> { |
| 223 | move |acc, x| { |
| 224 | n -= 1; |
| 225 | let r = fold(acc, x); |
| 226 | if n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) } |
| 227 | } |
| 228 | } |
| 229 | |
| 230 | let n = self.len(); |
| 231 | if n == 0 { try { init } } else { self.iter.try_rfold(init, check(n, fold)).into_try() } |
| 232 | } |
| 233 | |
| 234 | impl_fold_via_try_fold! { rfold -> try_rfold } |
| 235 | |
| 236 | #[inline ] |
| 237 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
| 238 | let min = crate::cmp::min(self.len(), n); |
| 239 | let rem = self.iter.advance_back_by(min); |
| 240 | assert!(rem.is_ok(), "ExactSizeIterator contract violation" ); |
| 241 | NonZero::new(n - min).map_or(Ok(()), Err) |
| 242 | } |
| 243 | } |
| 244 | |
| 245 | #[stable (feature = "fused" , since = "1.26.0" )] |
| 246 | impl<I> FusedIterator for Skip<I> where I: FusedIterator {} |
| 247 | |
| 248 | #[unstable (issue = "none" , feature = "trusted_fused" )] |
| 249 | unsafe impl<I: TrustedFused> TrustedFused for Skip<I> {} |
| 250 | |
| 251 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
| 252 | unsafe impl<I> SourceIter for Skip<I> |
| 253 | where |
| 254 | I: SourceIter, |
| 255 | { |
| 256 | type Source = I::Source; |
| 257 | |
| 258 | #[inline ] |
| 259 | unsafe fn as_inner(&mut self) -> &mut I::Source { |
| 260 | // SAFETY: unsafe function forwarding to unsafe function with the same requirements |
| 261 | unsafe { SourceIter::as_inner(&mut self.iter) } |
| 262 | } |
| 263 | } |
| 264 | |
| 265 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
| 266 | unsafe impl<I: InPlaceIterable> InPlaceIterable for Skip<I> { |
| 267 | const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY; |
| 268 | const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY; |
| 269 | } |
| 270 | |
| 271 | #[doc (hidden)] |
| 272 | #[unstable (feature = "trusted_random_access" , issue = "none" )] |
| 273 | unsafe impl<I> TrustedRandomAccess for Skip<I> where I: TrustedRandomAccess {} |
| 274 | |
| 275 | #[doc (hidden)] |
| 276 | #[unstable (feature = "trusted_random_access" , issue = "none" )] |
| 277 | unsafe impl<I> TrustedRandomAccessNoCoerce for Skip<I> |
| 278 | where |
| 279 | I: TrustedRandomAccessNoCoerce, |
| 280 | { |
| 281 | const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; |
| 282 | } |
| 283 | |
| 284 | // SAFETY: This adapter is shortening. TrustedLen requires the upper bound to be calculated correctly. |
| 285 | // These requirements can only be satisfied when the upper bound of the inner iterator's upper |
| 286 | // bound is never `None`. I: TrustedRandomAccess happens to provide this guarantee while |
| 287 | // I: TrustedLen would not. |
| 288 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
| 289 | unsafe impl<I> TrustedLen for Skip<I> where I: Iterator + TrustedRandomAccess {} |
| 290 | |