1 | use crate::intrinsics; |
2 | use crate::iter::{TrustedLen, TrustedRandomAccess, from_fn}; |
3 | use crate::num::NonZero; |
4 | use crate::ops::{Range, Try}; |
5 | |
6 | /// An iterator for stepping iterators by a custom amount. |
7 | /// |
8 | /// This `struct` is created by the [`step_by`] method on [`Iterator`]. See |
9 | /// its documentation for more. |
10 | /// |
11 | /// [`step_by`]: Iterator::step_by |
12 | /// [`Iterator`]: trait.Iterator.html |
13 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
14 | #[stable (feature = "iterator_step_by" , since = "1.28.0" )] |
15 | #[derive (Clone, Debug)] |
16 | pub struct StepBy<I> { |
17 | /// This field is guaranteed to be preprocessed by the specialized `SpecRangeSetup::setup` |
18 | /// in the constructor. |
19 | /// For most iterators that processing is a no-op, but for Range<{integer}> types it is lossy |
20 | /// which means the inner iterator cannot be returned to user code. |
21 | /// Additionally this type-dependent preprocessing means specialized implementations |
22 | /// cannot be used interchangeably. |
23 | iter: I, |
24 | /// This field is `step - 1`, aka the correct amount to pass to `nth` when iterating. |
25 | /// It MUST NOT be `usize::MAX`, as `unsafe` code depends on being able to add one |
26 | /// without the risk of overflow. (This is important so that length calculations |
27 | /// don't need to check for division-by-zero, for example.) |
28 | step_minus_one: usize, |
29 | first_take: bool, |
30 | } |
31 | |
32 | impl<I> StepBy<I> { |
33 | #[inline ] |
34 | pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> { |
35 | assert!(step != 0); |
36 | let iter: I = <I as SpecRangeSetup<I>>::setup(inner:iter, step); |
37 | StepBy { iter, step_minus_one: step - 1, first_take: true } |
38 | } |
39 | |
40 | /// The `step` that was originally passed to `Iterator::step_by(step)`, |
41 | /// aka `self.step_minus_one + 1`. |
42 | #[inline ] |
43 | fn original_step(&self) -> NonZero<usize> { |
44 | // SAFETY: By type invariant, `step_minus_one` cannot be `MAX`, which |
45 | // means the addition cannot overflow and the result cannot be zero. |
46 | unsafe { NonZero::new_unchecked(intrinsics::unchecked_add(self.step_minus_one, y:1)) } |
47 | } |
48 | } |
49 | |
50 | #[stable (feature = "iterator_step_by" , since = "1.28.0" )] |
51 | impl<I> Iterator for StepBy<I> |
52 | where |
53 | I: Iterator, |
54 | { |
55 | type Item = I::Item; |
56 | |
57 | #[inline ] |
58 | fn next(&mut self) -> Option<Self::Item> { |
59 | self.spec_next() |
60 | } |
61 | |
62 | #[inline ] |
63 | fn size_hint(&self) -> (usize, Option<usize>) { |
64 | self.spec_size_hint() |
65 | } |
66 | |
67 | #[inline ] |
68 | fn nth(&mut self, n: usize) -> Option<Self::Item> { |
69 | self.spec_nth(n) |
70 | } |
71 | |
72 | fn try_fold<Acc, F, R>(&mut self, acc: Acc, f: F) -> R |
73 | where |
74 | F: FnMut(Acc, Self::Item) -> R, |
75 | R: Try<Output = Acc>, |
76 | { |
77 | self.spec_try_fold(acc, f) |
78 | } |
79 | |
80 | #[inline ] |
81 | fn fold<Acc, F>(self, acc: Acc, f: F) -> Acc |
82 | where |
83 | F: FnMut(Acc, Self::Item) -> Acc, |
84 | { |
85 | self.spec_fold(acc, f) |
86 | } |
87 | } |
88 | |
89 | impl<I> StepBy<I> |
90 | where |
91 | I: ExactSizeIterator, |
92 | { |
93 | // The zero-based index starting from the end of the iterator of the |
94 | // last element. Used in the `DoubleEndedIterator` implementation. |
95 | fn next_back_index(&self) -> usize { |
96 | let rem: usize = self.iter.len() % self.original_step(); |
97 | if self.first_take { if rem == 0 { self.step_minus_one } else { rem - 1 } } else { rem } |
98 | } |
99 | } |
100 | |
101 | #[stable (feature = "double_ended_step_by_iterator" , since = "1.38.0" )] |
102 | impl<I> DoubleEndedIterator for StepBy<I> |
103 | where |
104 | I: DoubleEndedIterator + ExactSizeIterator, |
105 | { |
106 | #[inline ] |
107 | fn next_back(&mut self) -> Option<Self::Item> { |
108 | self.spec_next_back() |
109 | } |
110 | |
111 | #[inline ] |
112 | fn nth_back(&mut self, n: usize) -> Option<Self::Item> { |
113 | self.spec_nth_back(n) |
114 | } |
115 | |
116 | fn try_rfold<Acc, F, R>(&mut self, init: Acc, f: F) -> R |
117 | where |
118 | F: FnMut(Acc, Self::Item) -> R, |
119 | R: Try<Output = Acc>, |
120 | { |
121 | self.spec_try_rfold(init, f) |
122 | } |
123 | |
124 | #[inline ] |
125 | fn rfold<Acc, F>(self, init: Acc, f: F) -> Acc |
126 | where |
127 | Self: Sized, |
128 | F: FnMut(Acc, Self::Item) -> Acc, |
129 | { |
130 | self.spec_rfold(init, f) |
131 | } |
132 | } |
133 | |
134 | // StepBy can only make the iterator shorter, so the len will still fit. |
135 | #[stable (feature = "iterator_step_by" , since = "1.28.0" )] |
136 | impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {} |
137 | |
138 | // SAFETY: This adapter is shortening. TrustedLen requires the upper bound to be calculated correctly. |
139 | // These requirements can only be satisfied when the upper bound of the inner iterator's upper |
140 | // bound is never `None`. I: TrustedRandomAccess happens to provide this guarantee while |
141 | // I: TrustedLen would not. |
142 | // This also covers the Range specializations since the ranges also implement TRA |
143 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
144 | unsafe impl<I> TrustedLen for StepBy<I> where I: Iterator + TrustedRandomAccess {} |
145 | |
146 | trait SpecRangeSetup<T> { |
147 | fn setup(inner: T, step: usize) -> T; |
148 | } |
149 | |
150 | impl<T> SpecRangeSetup<T> for T { |
151 | #[inline ] |
152 | default fn setup(inner: T, _step: usize) -> T { |
153 | inner |
154 | } |
155 | } |
156 | |
157 | /// Specialization trait to optimize `StepBy<Range<{integer}>>` iteration. |
158 | /// |
159 | /// # Safety |
160 | /// |
161 | /// Technically this is safe to implement (look ma, no unsafe!), but in reality |
162 | /// a lot of unsafe code relies on ranges over integers being correct. |
163 | /// |
164 | /// For correctness *all* public StepBy methods must be specialized |
165 | /// because `setup` drastically alters the meaning of the struct fields so that mixing |
166 | /// different implementations would lead to incorrect results. |
167 | unsafe trait StepByImpl<I> { |
168 | type Item; |
169 | |
170 | fn spec_next(&mut self) -> Option<Self::Item>; |
171 | |
172 | fn spec_size_hint(&self) -> (usize, Option<usize>); |
173 | |
174 | fn spec_nth(&mut self, n: usize) -> Option<Self::Item>; |
175 | |
176 | fn spec_try_fold<Acc, F, R>(&mut self, acc: Acc, f: F) -> R |
177 | where |
178 | F: FnMut(Acc, Self::Item) -> R, |
179 | R: Try<Output = Acc>; |
180 | |
181 | fn spec_fold<Acc, F>(self, acc: Acc, f: F) -> Acc |
182 | where |
183 | F: FnMut(Acc, Self::Item) -> Acc; |
184 | } |
185 | |
186 | /// Specialization trait for double-ended iteration. |
187 | /// |
188 | /// See also: `StepByImpl` |
189 | /// |
190 | /// # Safety |
191 | /// |
192 | /// The specializations must be implemented together with `StepByImpl` |
193 | /// where applicable. I.e. if `StepBy` does support backwards iteration |
194 | /// for a given iterator and that is specialized for forward iteration then |
195 | /// it must also be specialized for backwards iteration. |
196 | unsafe trait StepByBackImpl<I> { |
197 | type Item; |
198 | |
199 | fn spec_next_back(&mut self) -> Option<Self::Item> |
200 | where |
201 | I: DoubleEndedIterator + ExactSizeIterator; |
202 | |
203 | fn spec_nth_back(&mut self, n: usize) -> Option<Self::Item> |
204 | where |
205 | I: DoubleEndedIterator + ExactSizeIterator; |
206 | |
207 | fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, f: F) -> R |
208 | where |
209 | I: DoubleEndedIterator + ExactSizeIterator, |
210 | F: FnMut(Acc, Self::Item) -> R, |
211 | R: Try<Output = Acc>; |
212 | |
213 | fn spec_rfold<Acc, F>(self, init: Acc, f: F) -> Acc |
214 | where |
215 | I: DoubleEndedIterator + ExactSizeIterator, |
216 | F: FnMut(Acc, Self::Item) -> Acc; |
217 | } |
218 | |
219 | unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> { |
220 | type Item = I::Item; |
221 | |
222 | #[inline ] |
223 | default fn spec_next(&mut self) -> Option<I::Item> { |
224 | let step_size = if self.first_take { 0 } else { self.step_minus_one }; |
225 | self.first_take = false; |
226 | self.iter.nth(step_size) |
227 | } |
228 | |
229 | #[inline ] |
230 | default fn spec_size_hint(&self) -> (usize, Option<usize>) { |
231 | #[inline ] |
232 | fn first_size(step: NonZero<usize>) -> impl Fn(usize) -> usize { |
233 | move |n| if n == 0 { 0 } else { 1 + (n - 1) / step } |
234 | } |
235 | |
236 | #[inline ] |
237 | fn other_size(step: NonZero<usize>) -> impl Fn(usize) -> usize { |
238 | move |n| n / step |
239 | } |
240 | |
241 | let (low, high) = self.iter.size_hint(); |
242 | |
243 | if self.first_take { |
244 | let f = first_size(self.original_step()); |
245 | (f(low), high.map(f)) |
246 | } else { |
247 | let f = other_size(self.original_step()); |
248 | (f(low), high.map(f)) |
249 | } |
250 | } |
251 | |
252 | #[inline ] |
253 | default fn spec_nth(&mut self, mut n: usize) -> Option<I::Item> { |
254 | if self.first_take { |
255 | self.first_take = false; |
256 | let first = self.iter.next(); |
257 | if n == 0 { |
258 | return first; |
259 | } |
260 | n -= 1; |
261 | } |
262 | // n and self.step_minus_one are indices, we need to add 1 to get the amount of elements |
263 | // When calling `.nth`, we need to subtract 1 again to convert back to an index |
264 | let mut step = self.original_step().get(); |
265 | // n + 1 could overflow |
266 | // thus, if n is usize::MAX, instead of adding one, we call .nth(step) |
267 | if n == usize::MAX { |
268 | self.iter.nth(step - 1); |
269 | } else { |
270 | n += 1; |
271 | } |
272 | |
273 | // overflow handling |
274 | loop { |
275 | let mul = n.checked_mul(step); |
276 | { |
277 | if intrinsics::likely(mul.is_some()) { |
278 | return self.iter.nth(mul.unwrap() - 1); |
279 | } |
280 | } |
281 | let div_n = usize::MAX / n; |
282 | let div_step = usize::MAX / step; |
283 | let nth_n = div_n * n; |
284 | let nth_step = div_step * step; |
285 | let nth = if nth_n > nth_step { |
286 | step -= div_n; |
287 | nth_n |
288 | } else { |
289 | n -= div_step; |
290 | nth_step |
291 | }; |
292 | self.iter.nth(nth - 1); |
293 | } |
294 | } |
295 | |
296 | default fn spec_try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R |
297 | where |
298 | F: FnMut(Acc, Self::Item) -> R, |
299 | R: Try<Output = Acc>, |
300 | { |
301 | #[inline ] |
302 | fn nth<I: Iterator>( |
303 | iter: &mut I, |
304 | step_minus_one: usize, |
305 | ) -> impl FnMut() -> Option<I::Item> + '_ { |
306 | move || iter.nth(step_minus_one) |
307 | } |
308 | |
309 | if self.first_take { |
310 | self.first_take = false; |
311 | match self.iter.next() { |
312 | None => return try { acc }, |
313 | Some(x) => acc = f(acc, x)?, |
314 | } |
315 | } |
316 | from_fn(nth(&mut self.iter, self.step_minus_one)).try_fold(acc, f) |
317 | } |
318 | |
319 | default fn spec_fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc |
320 | where |
321 | F: FnMut(Acc, Self::Item) -> Acc, |
322 | { |
323 | #[inline ] |
324 | fn nth<I: Iterator>( |
325 | iter: &mut I, |
326 | step_minus_one: usize, |
327 | ) -> impl FnMut() -> Option<I::Item> + '_ { |
328 | move || iter.nth(step_minus_one) |
329 | } |
330 | |
331 | if self.first_take { |
332 | self.first_take = false; |
333 | match self.iter.next() { |
334 | None => return acc, |
335 | Some(x) => acc = f(acc, x), |
336 | } |
337 | } |
338 | from_fn(nth(&mut self.iter, self.step_minus_one)).fold(acc, f) |
339 | } |
340 | } |
341 | |
342 | unsafe impl<I: DoubleEndedIterator + ExactSizeIterator> StepByBackImpl<I> for StepBy<I> { |
343 | type Item = I::Item; |
344 | |
345 | #[inline ] |
346 | default fn spec_next_back(&mut self) -> Option<Self::Item> { |
347 | self.iter.nth_back(self.next_back_index()) |
348 | } |
349 | |
350 | #[inline ] |
351 | default fn spec_nth_back(&mut self, n: usize) -> Option<I::Item> { |
352 | // `self.iter.nth_back(usize::MAX)` does the right thing here when `n` |
353 | // is out of bounds because the length of `self.iter` does not exceed |
354 | // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is |
355 | // zero-indexed |
356 | let n = n.saturating_mul(self.original_step().get()).saturating_add(self.next_back_index()); |
357 | self.iter.nth_back(n) |
358 | } |
359 | |
360 | default fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R |
361 | where |
362 | F: FnMut(Acc, Self::Item) -> R, |
363 | R: Try<Output = Acc>, |
364 | { |
365 | #[inline ] |
366 | fn nth_back<I: DoubleEndedIterator>( |
367 | iter: &mut I, |
368 | step_minus_one: usize, |
369 | ) -> impl FnMut() -> Option<I::Item> + '_ { |
370 | move || iter.nth_back(step_minus_one) |
371 | } |
372 | |
373 | match self.next_back() { |
374 | None => try { init }, |
375 | Some(x) => { |
376 | let acc = f(init, x)?; |
377 | from_fn(nth_back(&mut self.iter, self.step_minus_one)).try_fold(acc, f) |
378 | } |
379 | } |
380 | } |
381 | |
382 | #[inline ] |
383 | default fn spec_rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc |
384 | where |
385 | Self: Sized, |
386 | F: FnMut(Acc, I::Item) -> Acc, |
387 | { |
388 | #[inline ] |
389 | fn nth_back<I: DoubleEndedIterator>( |
390 | iter: &mut I, |
391 | step_minus_one: usize, |
392 | ) -> impl FnMut() -> Option<I::Item> + '_ { |
393 | move || iter.nth_back(step_minus_one) |
394 | } |
395 | |
396 | match self.next_back() { |
397 | None => init, |
398 | Some(x) => { |
399 | let acc = f(init, x); |
400 | from_fn(nth_back(&mut self.iter, self.step_minus_one)).fold(acc, f) |
401 | } |
402 | } |
403 | } |
404 | } |
405 | |
406 | /// For these implementations, `SpecRangeSetup` calculates the number |
407 | /// of iterations that will be needed and stores that in `iter.end`. |
408 | /// |
409 | /// The various iterator implementations then rely on that to not need |
410 | /// overflow checking, letting loops just be counted instead. |
411 | /// |
412 | /// These only work for unsigned types, and will need to be reworked |
413 | /// if you want to use it to specialize on signed types. |
414 | /// |
415 | /// Currently these are only implemented for integers up to `usize` due to |
416 | /// correctness issues around `ExactSizeIterator` impls on 16bit platforms. |
417 | /// And since `ExactSizeIterator` is a prerequisite for backwards iteration |
418 | /// and we must consistently specialize backwards and forwards iteration |
419 | /// that makes the situation complicated enough that it's not covered |
420 | /// for now. |
421 | macro_rules! spec_int_ranges { |
422 | ($($t:ty)*) => ($( |
423 | |
424 | const _: () = assert!(usize::BITS >= <$t>::BITS); |
425 | |
426 | impl SpecRangeSetup<Range<$t>> for Range<$t> { |
427 | #[inline] |
428 | fn setup(mut r: Range<$t>, step: usize) -> Range<$t> { |
429 | let inner_len = r.size_hint().0; |
430 | // If step exceeds $t::MAX, then the count will be at most 1 and |
431 | // thus always fit into $t. |
432 | let yield_count = inner_len.div_ceil(step); |
433 | // Turn the range end into an iteration counter |
434 | r.end = yield_count as $t; |
435 | r |
436 | } |
437 | } |
438 | |
439 | unsafe impl StepByImpl<Range<$t>> for StepBy<Range<$t>> { |
440 | #[inline] |
441 | fn spec_next(&mut self) -> Option<$t> { |
442 | // if a step size larger than the type has been specified fall back to |
443 | // t::MAX, in which case remaining will be at most 1. |
444 | let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX); |
445 | let remaining = self.iter.end; |
446 | if remaining > 0 { |
447 | let val = self.iter.start; |
448 | // this can only overflow during the last step, after which the value |
449 | // will not be used |
450 | self.iter.start = val.wrapping_add(step); |
451 | self.iter.end = remaining - 1; |
452 | Some(val) |
453 | } else { |
454 | None |
455 | } |
456 | } |
457 | |
458 | #[inline] |
459 | fn spec_size_hint(&self) -> (usize, Option<usize>) { |
460 | let remaining = self.iter.end as usize; |
461 | (remaining, Some(remaining)) |
462 | } |
463 | |
464 | // The methods below are all copied from the Iterator trait default impls. |
465 | // We have to repeat them here so that the specialization overrides the StepByImpl defaults |
466 | |
467 | #[inline] |
468 | fn spec_nth(&mut self, n: usize) -> Option<Self::Item> { |
469 | self.advance_by(n).ok()?; |
470 | self.next() |
471 | } |
472 | |
473 | #[inline] |
474 | fn spec_try_fold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R |
475 | where |
476 | F: FnMut(Acc, Self::Item) -> R, |
477 | R: Try<Output = Acc> |
478 | { |
479 | let mut accum = init; |
480 | while let Some(x) = self.next() { |
481 | accum = f(accum, x)?; |
482 | } |
483 | try { accum } |
484 | } |
485 | |
486 | #[inline] |
487 | fn spec_fold<Acc, F>(self, init: Acc, mut f: F) -> Acc |
488 | where |
489 | F: FnMut(Acc, Self::Item) -> Acc |
490 | { |
491 | // if a step size larger than the type has been specified fall back to |
492 | // t::MAX, in which case remaining will be at most 1. |
493 | let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX); |
494 | let remaining = self.iter.end; |
495 | let mut acc = init; |
496 | let mut val = self.iter.start; |
497 | for _ in 0..remaining { |
498 | acc = f(acc, val); |
499 | // this can only overflow during the last step, after which the value |
500 | // will no longer be used |
501 | val = val.wrapping_add(step); |
502 | } |
503 | acc |
504 | } |
505 | } |
506 | )*) |
507 | } |
508 | |
509 | macro_rules! spec_int_ranges_r { |
510 | ($($t:ty)*) => ($( |
511 | const _: () = assert!(usize::BITS >= <$t>::BITS); |
512 | |
513 | unsafe impl StepByBackImpl<Range<$t>> for StepBy<Range<$t>> { |
514 | |
515 | #[inline] |
516 | fn spec_next_back(&mut self) -> Option<Self::Item> { |
517 | let step = self.original_step().get() as $t; |
518 | let remaining = self.iter.end; |
519 | if remaining > 0 { |
520 | let start = self.iter.start; |
521 | self.iter.end = remaining - 1; |
522 | Some(start + step * (remaining - 1)) |
523 | } else { |
524 | None |
525 | } |
526 | } |
527 | |
528 | // The methods below are all copied from the Iterator trait default impls. |
529 | // We have to repeat them here so that the specialization overrides the StepByImplBack defaults |
530 | |
531 | #[inline] |
532 | fn spec_nth_back(&mut self, n: usize) -> Option<Self::Item> { |
533 | if self.advance_back_by(n).is_err() { |
534 | return None; |
535 | } |
536 | self.next_back() |
537 | } |
538 | |
539 | #[inline] |
540 | fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R |
541 | where |
542 | F: FnMut(Acc, Self::Item) -> R, |
543 | R: Try<Output = Acc> |
544 | { |
545 | let mut accum = init; |
546 | while let Some(x) = self.next_back() { |
547 | accum = f(accum, x)?; |
548 | } |
549 | try { accum } |
550 | } |
551 | |
552 | #[inline] |
553 | fn spec_rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc |
554 | where |
555 | F: FnMut(Acc, Self::Item) -> Acc |
556 | { |
557 | let mut accum = init; |
558 | while let Some(x) = self.next_back() { |
559 | accum = f(accum, x); |
560 | } |
561 | accum |
562 | } |
563 | } |
564 | )*) |
565 | } |
566 | |
567 | #[cfg (target_pointer_width = "64" )] |
568 | spec_int_ranges!(u8 u16 u32 u64 usize); |
569 | // DoubleEndedIterator requires ExactSizeIterator, which isn't implemented for Range<u64> |
570 | #[cfg (target_pointer_width = "64" )] |
571 | spec_int_ranges_r!(u8 u16 u32 usize); |
572 | |
573 | #[cfg (target_pointer_width = "32" )] |
574 | spec_int_ranges!(u8 u16 u32 usize); |
575 | #[cfg (target_pointer_width = "32" )] |
576 | spec_int_ranges_r!(u8 u16 u32 usize); |
577 | |
578 | #[cfg (target_pointer_width = "16" )] |
579 | spec_int_ranges!(u8 u16 usize); |
580 | #[cfg (target_pointer_width = "16" )] |
581 | spec_int_ranges_r!(u8 u16 usize); |
582 | |