1use crate::intrinsics;
2use crate::iter::{TrustedLen, TrustedRandomAccess, from_fn};
3use crate::num::NonZero;
4use crate::ops::{Range, Try};
5
6#[must_use = "iterators are lazy and do nothing unless consumed"]
14#[stable(feature = "iterator_step_by", since = "1.28.0")]
15#[derive(Clone, Debug)]
16pub struct StepBy<I> {
17    iter: I,
24    step_minus_one: usize,
29    first_take: bool,
30}
31
32impl<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 as SpecRangeSetup<I>>::setup(iter, step);
37        StepBy { iter, step_minus_one: step - 1, first_take: true }
38    }
39
40    #[inline]
43    fn original_step(&self) -> NonZero<usize> {
44        unsafe { NonZero::new_unchecked(intrinsics::unchecked_add(self.step_minus_one, 1)) }
47    }
48}
49
50#[stable(feature = "iterator_step_by", since = "1.28.0")]
51impl<I> Iterator for StepBy<I>
52where
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
89impl<I> StepBy<I>
90where
91    I: ExactSizeIterator,
92{
93    fn next_back_index(&self) -> usize {
96        let rem = 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")]
102impl<I> DoubleEndedIterator for StepBy<I>
103where
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#[stable(feature = "iterator_step_by", since = "1.28.0")]
136impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
137
138#[unstable(feature = "trusted_len", issue = "37572")]
144unsafe impl<I> TrustedLen for StepBy<I> where I: Iterator + TrustedRandomAccess {}
145
146trait SpecRangeSetup<T> {
147    fn setup(inner: T, step: usize) -> T;
148}
149
150impl<T> SpecRangeSetup<T> for T {
151    #[inline]
152    default fn setup(inner: T, _step: usize) -> T {
153        inner
154    }
155}
156
157unsafe 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
186unsafe 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
219unsafe 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        let mut step = self.original_step().get();
265        if n == usize::MAX {
268            self.iter.nth(step - 1);
269        } else {
270            n += 1;
271        }
272
273        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
342unsafe 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        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
406macro_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                let yield_count = inner_len.div_ceil(step);
433                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                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                    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            #[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                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                    val = val.wrapping_add(step);
502                }
503                acc
504            }
505        }
506    )*)
507}
508
509macro_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            #[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")]
568spec_int_ranges!(u8 u16 u32 u64 usize);
569#[cfg(target_pointer_width = "64")]
571spec_int_ranges_r!(u8 u16 u32 usize);
572
573#[cfg(target_pointer_width = "32")]
574spec_int_ranges!(u8 u16 u32 usize);
575#[cfg(target_pointer_width = "32")]
576spec_int_ranges_r!(u8 u16 u32 usize);
577
578#[cfg(target_pointer_width = "16")]
579spec_int_ranges!(u8 u16 usize);
580#[cfg(target_pointer_width = "16")]
581spec_int_ranges_r!(u8 u16 usize);