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

core/iter/adapters/
step_by.rs
