From dcc9c46efc68eb50069fd179182b5cf3b1c545a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E6=99=A8=20=28Leo=20Cheng=29?= Date: Sat, 25 Jul 2026 21:25:46 +0800 Subject: [PATCH] iter: extend step_by specialization to cover StepBy> MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 林晨 (Leo Cheng) --- library/core/src/iter/adapters/step_by.rs | 110 +++++++++++++----- library/core/src/range/iter.rs | 2 +- .../coretests/tests/iter/adapters/step_by.rs | 38 ++++++ 3 files changed, 122 insertions(+), 28 deletions(-) diff --git a/library/core/src/iter/adapters/step_by.rs b/library/core/src/iter/adapters/step_by.rs index 2d0f210420317..4c70c9dc84b3b 100644 --- a/library/core/src/iter/adapters/step_by.rs +++ b/library/core/src/iter/adapters/step_by.rs @@ -2,6 +2,7 @@ use crate::intrinsics; use crate::iter::{TrustedLen, TrustedRandomAccess, from_fn}; use crate::num::NonZero; use crate::ops::{Range, Try}; +use crate::range::RangeIter; /// An iterator for stepping iterators by a custom amount. /// @@ -418,37 +419,71 @@ unsafe impl StepByBackImpl for St /// and we must consistently specialize backwards and forwards iteration /// that makes the situation complicated enough that it's not covered /// for now. +/// +/// After `SpecRangeSetup::setup`, both `Range` and its new-range wrapper +/// `RangeIter` carry the cursor and countdown in the same underlying legacy +/// `Range`. This accessor exposes that shared range so one specialization can +/// serve both: it is an identity for `Range` and unwraps the newtype for +/// `RangeIter`, so it compiles away. +trait AsLegacyRange { + fn as_legacy_range(&self) -> &Range; + fn as_legacy_range_mut(&mut self) -> &mut Range; +} + +impl AsLegacyRange for Range { + #[inline] + fn as_legacy_range(&self) -> &Range { + self + } + #[inline] + fn as_legacy_range_mut(&mut self) -> &mut Range { + self + } +} + +impl AsLegacyRange for RangeIter { + #[inline] + fn as_legacy_range(&self) -> &Range { + &self.0 + } + #[inline] + fn as_legacy_range_mut(&mut self) -> &mut Range { + &mut self.0 + } +} + macro_rules! spec_int_ranges { - ($($t:ty)*) => ($( + ($ctor:ident; $($t:ty)*) => ($( const _: () = assert!(usize::BITS >= <$t>::BITS); - impl SpecRangeSetup> for Range<$t> { + impl SpecRangeSetup<$ctor<$t>> for $ctor<$t> { #[inline] - fn setup(mut r: Range<$t>, step: usize) -> Range<$t> { + fn setup(mut r: $ctor<$t>, step: usize) -> $ctor<$t> { let inner_len = r.size_hint().0; // If step exceeds $t::MAX, then the count will be at most 1 and // thus always fit into $t. let yield_count = inner_len.div_ceil(step); // Turn the range end into an iteration counter - r.end = yield_count as $t; + r.as_legacy_range_mut().end = yield_count as $t; r } } - unsafe impl StepByImpl> for StepBy> { + unsafe impl StepByImpl<$ctor<$t>> for StepBy<$ctor<$t>> { #[inline] fn spec_next(&mut self) -> Option<$t> { // if a step size larger than the type has been specified fall back to // t::MAX, in which case remaining will be at most 1. let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX); - let remaining = self.iter.end; + let r = self.iter.as_legacy_range_mut(); + let remaining = r.end; if remaining > 0 { - let val = self.iter.start; + let val = r.start; // this can only overflow during the last step, after which the value // will not be used - self.iter.start = val.wrapping_add(step); - self.iter.end = remaining - 1; + r.start = val.wrapping_add(step); + r.end = remaining - 1; Some(val) } else { None @@ -457,7 +492,7 @@ macro_rules! spec_int_ranges { #[inline] fn spec_size_hint(&self) -> (usize, Option) { - let remaining = self.iter.end as usize; + let remaining = self.iter.as_legacy_range().end as usize; (remaining, Some(remaining)) } @@ -491,9 +526,10 @@ macro_rules! spec_int_ranges { // if a step size larger than the type has been specified fall back to // t::MAX, in which case remaining will be at most 1. let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX); - let remaining = self.iter.end; + let r = self.iter.as_legacy_range(); + let remaining = r.end; let mut acc = init; - let mut val = self.iter.start; + let mut val = r.start; for _ in 0..remaining { acc = f(acc, val); // this can only overflow during the last step, after which the value @@ -507,18 +543,19 @@ macro_rules! spec_int_ranges { } macro_rules! spec_int_ranges_r { - ($($t:ty)*) => ($( + ($ctor:ident; $($t:ty)*) => ($( const _: () = assert!(usize::BITS >= <$t>::BITS); - unsafe impl StepByBackImpl> for StepBy> { + unsafe impl StepByBackImpl<$ctor<$t>> for StepBy<$ctor<$t>> { #[inline] fn spec_next_back(&mut self) -> Option { let step = self.original_step().get() as $t; - let remaining = self.iter.end; + let r = self.iter.as_legacy_range_mut(); + let remaining = r.end; if remaining > 0 { - let start = self.iter.start; - self.iter.end = remaining - 1; + let start = r.start; + r.end = remaining - 1; Some(start + step * (remaining - 1)) } else { None @@ -564,18 +601,37 @@ macro_rules! spec_int_ranges_r { )*) } +// The same specialization covers `Range<{integer}>` and the new-range iterator +// `RangeIter<{integer}>`, which wraps a `Range` (see `AsLegacyRange`). +// +// The backward (`_r`) specialization requires `ExactSizeIterator`. `RangeIter` +// implements it only for `usize`/`u8`/`u16` (see `range_exact_iter_impl!` in +// `range::iter`), narrower than `Range`, so `RangeIter`'s backward set omits +// `u32` even where `Range` includes it; `Range` is likewise omitted on +// 64-bit since its length can exceed `usize`. #[cfg(target_pointer_width = "64")] -spec_int_ranges!(u8 u16 u32 u64 usize); -// DoubleEndedIterator requires ExactSizeIterator, which isn't implemented for Range -#[cfg(target_pointer_width = "64")] -spec_int_ranges_r!(u8 u16 u32 usize); +mod step_by_spec { + use super::*; + spec_int_ranges!(Range; u8 u16 u32 u64 usize); + spec_int_ranges!(RangeIter; u8 u16 u32 u64 usize); + spec_int_ranges_r!(Range; u8 u16 u32 usize); + spec_int_ranges_r!(RangeIter; u8 u16 usize); +} #[cfg(target_pointer_width = "32")] -spec_int_ranges!(u8 u16 u32 usize); -#[cfg(target_pointer_width = "32")] -spec_int_ranges_r!(u8 u16 u32 usize); +mod step_by_spec { + use super::*; + spec_int_ranges!(Range; u8 u16 u32 usize); + spec_int_ranges!(RangeIter; u8 u16 u32 usize); + spec_int_ranges_r!(Range; u8 u16 u32 usize); + spec_int_ranges_r!(RangeIter; u8 u16 usize); +} #[cfg(target_pointer_width = "16")] -spec_int_ranges!(u8 u16 usize); -#[cfg(target_pointer_width = "16")] -spec_int_ranges_r!(u8 u16 usize); +mod step_by_spec { + use super::*; + spec_int_ranges!(Range; u8 u16 usize); + spec_int_ranges!(RangeIter; u8 u16 usize); + spec_int_ranges_r!(Range; u8 u16 usize); + spec_int_ranges_r!(RangeIter; u8 u16 usize); +} diff --git a/library/core/src/range/iter.rs b/library/core/src/range/iter.rs index 01b69554a0b1b..6011949476f6f 100644 --- a/library/core/src/range/iter.rs +++ b/library/core/src/range/iter.rs @@ -8,7 +8,7 @@ use crate::{intrinsics, mem}; /// By-value [`Range`] iterator. #[stable(feature = "new_range_api", since = "1.96.0")] #[derive(Debug, Clone)] -pub struct RangeIter(legacy::Range); +pub struct RangeIter(pub(crate) legacy::Range); impl RangeIter { #[unstable(feature = "new_range_remainder", issue = "154458")] diff --git a/library/coretests/tests/iter/adapters/step_by.rs b/library/coretests/tests/iter/adapters/step_by.rs index 6f3300e7a8820..1ebebb9691933 100644 --- a/library/coretests/tests/iter/adapters/step_by.rs +++ b/library/coretests/tests/iter/adapters/step_by.rs @@ -299,3 +299,41 @@ fn test_step_by_fold_range_specialization() { assert_eq!(r.sum::(), usize::MAX - 1); }); } + +#[test] +fn test_step_by_new_range_iter() { + use core::range::Range as NewRange; + + // forward iteration + let v: Vec = NewRange::from(0_u32..10).into_iter().step_by(3).collect(); + assert_eq!(v, [0, 3, 6, 9]); + + // size_hint + assert_eq!(NewRange::from(0_u32..10).into_iter().step_by(3).size_hint(), (4, Some(4))); + assert_eq!(NewRange::from(0_u32..9).into_iter().step_by(3).size_hint(), (3, Some(3))); + + // nth + assert_eq!(NewRange::from(0_u32..20).into_iter().step_by(5).nth(2), Some(10)); + + // empty range + assert_eq!(NewRange::from(5_u32..5).into_iter().step_by(1).next(), None); + + // step larger than range + assert_eq!(NewRange::from(0_u32..3).into_iter().step_by(10).collect::>(), [0]); + + // backward iteration (usize has ExactSizeIterator) + let mut it = NewRange::from(0_usize..11).into_iter().step_by(3); + assert_eq!(it.next_back(), Some(9)); + assert_eq!(it.next_back(), Some(6)); + assert_eq!(it.next_back(), Some(3)); + assert_eq!(it.next_back(), Some(0)); + assert_eq!(it.next_back(), None); + + // interleaved forward and backward + let mut it = NewRange::from(0_usize..16).into_iter().step_by(5); + assert_eq!(it.next(), Some(0)); + assert_eq!(it.next_back(), Some(15)); + assert_eq!(it.next(), Some(5)); + assert_eq!(it.next_back(), Some(10)); + assert_eq!(it.next(), None); +}