Skip to content

Commit fd15168

Browse files
committed
chore : Range has also been modified to be the same as the existing step_by.
1 parent 485ec3f commit fd15168

3 files changed

Lines changed: 147 additions & 1 deletion

File tree

library/core/src/iter/adapters/step_by.rs

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use crate::intrinsics;
22
use crate::iter::{TrustedLen, TrustedRandomAccess, from_fn};
33
use crate::num::NonZero;
44
use crate::ops::{Range, Try};
5+
use crate::range::RangeIter;
56

67
/// An iterator for stepping iterators by a custom amount.
78
///
@@ -503,6 +504,87 @@ macro_rules! spec_int_ranges {
503504
acc
504505
}
505506
}
507+
508+
impl SpecRangeSetup<RangeIter<$t>> for RangeIter<$t> {
509+
#[inline]
510+
fn setup(mut r: RangeIter<$t>, step: usize) -> RangeIter<$t> {
511+
let inner_len = r.size_hint().0;
512+
// If step exceeds $t::MAX, then the count will be at most 1 and
513+
// thus always fit into $t.
514+
let yield_count = inner_len.div_ceil(step);
515+
// Turn the range end into an iteration counter
516+
r.0.end = yield_count as $t;
517+
r
518+
}
519+
}
520+
521+
unsafe impl StepByImpl<RangeIter<$t>> for StepBy<RangeIter<$t>> {
522+
#[inline]
523+
fn spec_next(&mut self) -> Option<$t> {
524+
// if a step size larger than the type has been specified fall back to
525+
// t::MAX, in which case remaining will be at most 1.
526+
let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX);
527+
let remaining = self.iter.0.end;
528+
if remaining > 0 {
529+
let val = self.iter.0.start;
530+
// this can only overflow during the last step, after which the value
531+
// will not be used
532+
self.iter.0.start = val.wrapping_add(step);
533+
self.iter.0.end = remaining - 1;
534+
Some(val)
535+
} else {
536+
None
537+
}
538+
}
539+
540+
#[inline]
541+
fn spec_size_hint(&self) -> (usize, Option<usize>) {
542+
let remaining = self.iter.0.end as usize;
543+
(remaining, Some(remaining))
544+
}
545+
546+
// The methods below are all copied from the Iterator trait default impls.
547+
// We have to repeat them here so that the specialization overrides the StepByImpl defaults
548+
549+
#[inline]
550+
fn spec_nth(&mut self, n: usize) -> Option<Self::Item> {
551+
self.advance_by(n).ok()?;
552+
self.next()
553+
}
554+
555+
#[inline]
556+
fn spec_try_fold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
557+
where
558+
F: FnMut(Acc, Self::Item) -> R,
559+
R: Try<Output = Acc>
560+
{
561+
let mut accum = init;
562+
while let Some(x) = self.next() {
563+
accum = f(accum, x)?;
564+
}
565+
try { accum }
566+
}
567+
568+
#[inline]
569+
fn spec_fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
570+
where
571+
F: FnMut(Acc, Self::Item) -> Acc
572+
{
573+
// if a step size larger than the type has been specified fall back to
574+
// t::MAX, in which case remaining will be at most 1.
575+
let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX);
576+
let remaining = self.iter.0.end;
577+
let mut acc = init;
578+
let mut val = self.iter.0.start;
579+
for _ in 0..remaining {
580+
acc = f(acc, val);
581+
// this can only overflow during the last step, after which the value
582+
// will no longer be used
583+
val = val.wrapping_add(step);
584+
}
585+
acc
586+
}
587+
}
506588
)*)
507589
}
508590

@@ -564,11 +646,70 @@ macro_rules! spec_int_ranges_r {
564646
)*)
565647
}
566648

649+
macro_rules! spec_int_range_iters_r {
650+
($($t:ty)*) => ($(
651+
const _: () = assert!(usize::BITS >= <$t>::BITS);
652+
653+
unsafe impl StepByBackImpl<RangeIter<$t>> for StepBy<RangeIter<$t>> {
654+
655+
#[inline]
656+
fn spec_next_back(&mut self) -> Option<Self::Item> {
657+
let step = self.original_step().get() as $t;
658+
let remaining = self.iter.0.end;
659+
if remaining > 0 {
660+
let start = self.iter.0.start;
661+
self.iter.0.end = remaining - 1;
662+
Some(start + step * (remaining - 1))
663+
} else {
664+
None
665+
}
666+
}
667+
668+
// The methods below are all copied from the Iterator trait default impls.
669+
// We have to repeat them here so that the specialization overrides the StepByImplBack defaults
670+
671+
#[inline]
672+
fn spec_nth_back(&mut self, n: usize) -> Option<Self::Item> {
673+
if self.advance_back_by(n).is_err() {
674+
return None;
675+
}
676+
self.next_back()
677+
}
678+
679+
#[inline]
680+
fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
681+
where
682+
F: FnMut(Acc, Self::Item) -> R,
683+
R: Try<Output = Acc>
684+
{
685+
let mut accum = init;
686+
while let Some(x) = self.next_back() {
687+
accum = f(accum, x)?;
688+
}
689+
try { accum }
690+
}
691+
692+
#[inline]
693+
fn spec_rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
694+
where
695+
F: FnMut(Acc, Self::Item) -> Acc
696+
{
697+
let mut accum = init;
698+
while let Some(x) = self.next_back() {
699+
accum = f(accum, x);
700+
}
701+
accum
702+
}
703+
}
704+
)*)
705+
}
706+
567707
#[cfg(target_pointer_width = "64")]
568708
spec_int_ranges!(u8 u16 u32 u64 usize);
569709
// DoubleEndedIterator requires ExactSizeIterator, which isn't implemented for Range<u64>
570710
#[cfg(target_pointer_width = "64")]
571711
spec_int_ranges_r!(u8 u16 u32 usize);
712+
spec_int_range_iters_r!(u8 u16 usize);
572713

573714
#[cfg(target_pointer_width = "32")]
574715
spec_int_ranges!(u8 u16 u32 usize);

library/core/src/range/iter.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::{intrinsics, mem};
88
/// By-value [`Range`] iterator.
99
#[stable(feature = "new_range_api", since = "1.96.0")]
1010
#[derive(Debug, Clone)]
11-
pub struct RangeIter<A>(legacy::Range<A>);
11+
pub struct RangeIter<A>(pub(crate) legacy::Range<A>);
1212

1313
impl<A> RangeIter<A> {
1414
#[unstable(feature = "new_range_remainder", issue = "154458")]

library/coretests/tests/iter/adapters/step_by.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,11 @@ fn test_step_by_fold_range_specialization() {
294294
assert_eq!(r.sum::<usize>(), 6);
295295
});
296296

297+
t!((core::range::Range { start: 0usize, end: 5 }).into_iter().step_by(2), r, {
298+
assert_eq!(r.next(), Some(0));
299+
assert_eq!(r.sum::<usize>(), 6);
300+
});
301+
297302
t!((usize::MAX - 6..usize::MAX).step_by(5), r, {
298303
assert_eq!(r.next(), Some(usize::MAX - 6));
299304
assert_eq!(r.sum::<usize>(), usize::MAX - 1);

0 commit comments

Comments
 (0)