core/iter/adapters/
step_by.rs

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

core/iter/adapters/
step_by.rs
