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);