diff --git a/library/core/src/iter/range.rs b/library/core/src/iter/range.rs index 951bb5d0029f6..adc561879ebd3 100644 --- a/library/core/src/iter/range.rs +++ b/library/core/src/iter/range.rs @@ -15,6 +15,7 @@ macro_rules! unsafe_impl_trusted_step { )*}; } unsafe_impl_trusted_step![AsciiChar char i8 i16 i32 i64 i128 isize u8 u16 u32 u64 u128 usize Ipv4Addr Ipv6Addr]; +unsafe_impl_trusted_step![NonZero NonZero NonZero NonZero NonZero NonZero]; /// Objects that have a notion of *successor* and *predecessor* operations. /// @@ -255,10 +256,8 @@ macro_rules! step_identical_methods { macro_rules! step_integer_impls { { - narrower than or same width as usize: - $( [ $u_narrower:ident $i_narrower:ident ] ),+; - wider than usize: - $( [ $u_wider:ident $i_wider:ident ] ),+; + [ $( [ $u_narrower:ident $i_narrower:ident ] ),+ ] <= usize < + [ $( [ $u_wider:ident $i_wider:ident ] ),+ ] } => { $( #[allow(unreachable_patterns)] @@ -437,20 +436,138 @@ macro_rules! step_integer_impls { #[cfg(target_pointer_width = "64")] step_integer_impls! { - narrower than or same width as usize: [u8 i8], [u16 i16], [u32 i32], [u64 i64], [usize isize]; - wider than usize: [u128 i128]; + [ [u8 i8], [u16 i16], [u32 i32], [u64 i64], [usize isize] ] <= usize < [ [u128 i128] ] } #[cfg(target_pointer_width = "32")] step_integer_impls! { - narrower than or same width as usize: [u8 i8], [u16 i16], [u32 i32], [usize isize]; - wider than usize: [u64 i64], [u128 i128]; + [ [u8 i8], [u16 i16], [u32 i32], [usize isize] ] <= usize < [ [u64 i64], [u128 i128] ] } #[cfg(target_pointer_width = "16")] step_integer_impls! { - narrower than or same width as usize: [u8 i8], [u16 i16], [usize isize]; - wider than usize: [u32 i32], [u64 i64], [u128 i128]; + [ [u8 i8], [u16 i16], [usize isize] ] <= usize < [ [u32 i32], [u64 i64], [u128 i128] ] +} + +// These are still macro-generated because the integer literals resolve to different types. +macro_rules! step_nonzero_identical_methods { + ($int:ident) => { + #[inline] + unsafe fn forward_unchecked(start: Self, n: usize) -> Self { + // SAFETY: the caller has to guarantee that `start + n` doesn't overflow. + unsafe { Self::new_unchecked(start.get().unchecked_add(n as $int)) } + } + + #[inline] + unsafe fn backward_unchecked(start: Self, n: usize) -> Self { + // SAFETY: the caller has to guarantee that `start - n` doesn't overflow or hit zero. + unsafe { Self::new_unchecked(start.get().unchecked_sub(n as $int)) } + } + + #[inline] + #[allow(arithmetic_overflow)] + #[rustc_inherit_overflow_checks] + fn forward(start: Self, n: usize) -> Self { + // In debug builds, trigger a panic on overflow. + // This should optimize completely out in release builds. + if Self::forward_checked(start, n).is_none() { + let _ = $int::MAX + 1; + } + // Do saturating math (wrapping math causes UB if it wraps to Zero) + start.saturating_add(n as $int) + } + + #[inline] + #[allow(arithmetic_overflow)] + #[rustc_inherit_overflow_checks] + fn backward(start: Self, n: usize) -> Self { + // In debug builds, trigger a panic on overflow. + // This should optimize completely out in release builds. + if Self::backward_checked(start, n).is_none() { + let _ = $int::MIN - 1; + } + // Do saturating math (wrapping math causes UB if it wraps to Zero) + Self::new(start.get().saturating_sub(n as $int)).unwrap_or(Self::MIN) + } + + #[inline] + fn steps_between(start: &Self, end: &Self) -> (usize, Option) { + if *start <= *end { + #[allow(irrefutable_let_patterns, reason = "happens on usize or narrower")] + if let Ok(steps) = usize::try_from(end.get() - start.get()) { + (steps, Some(steps)) + } else { + (usize::MAX, None) + } + } else { + (0, None) + } + } + }; +} + +macro_rules! step_nonzero_impls { + { + [$( $narrower:ident ),+] <= usize < [$( $wider:ident ),+] + } => { + $( + #[allow(unreachable_patterns)] + #[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")] + impl Step for NonZero<$narrower> { + step_nonzero_identical_methods!($narrower); + + #[inline] + fn forward_checked(start: Self, n: usize) -> Option { + match $narrower::try_from(n) { + Ok(n) => start.checked_add(n), + Err(_) => None, // if n is out of range, `unsigned_start + n` is too + } + } + + #[inline] + fn backward_checked(start: Self, n: usize) -> Option { + match $narrower::try_from(n) { + // *_sub() is not implemented on NonZero + Ok(n) => start.get().checked_sub(n).and_then(Self::new), + Err(_) => None, // if n is out of range, `unsigned_start - n` is too + } + } + } + )+ + + $( + #[allow(unreachable_patterns)] + #[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")] + impl Step for NonZero<$wider> { + step_nonzero_identical_methods!($wider); + + #[inline] + fn forward_checked(start: Self, n: usize) -> Option { + start.checked_add(n as $wider) + } + + #[inline] + fn backward_checked(start: Self, n: usize) -> Option { + start.get().checked_sub(n as $wider).and_then(Self::new) + } + } + )+ + }; +} + +#[cfg(target_pointer_width = "64")] +step_nonzero_impls! { + [u8, u16, u32, u64, usize] <= usize < [u128] +} + +#[cfg(target_pointer_width = "32")] +step_nonzero_impls! { + [u8, u16, u32, usize] <= usize < [u64, u128] +} + +#[cfg(target_pointer_width = "16")] +step_nonzero_impls! { + [u8, u16, usize] <= usize < [u32, u64, u128] } #[unstable(feature = "step_trait", issue = "42168")] @@ -944,6 +1061,7 @@ impl Iterator for ops::Range { range_exact_iter_impl! { usize u8 u16 isize i8 i16 + NonZero NonZero NonZero // These are incorrect per the reasoning above, // but removing them would be a breaking change as they were stabilized in Rust 1.0.0. @@ -956,22 +1074,30 @@ range_exact_iter_impl! { unsafe_range_trusted_random_access_impl! { usize u8 u16 isize i8 i16 + NonZero NonZero NonZero } #[cfg(target_pointer_width = "32")] unsafe_range_trusted_random_access_impl! { u32 i32 + NonZero } #[cfg(target_pointer_width = "64")] unsafe_range_trusted_random_access_impl! { u32 i32 u64 i64 + NonZero + NonZero } range_incl_exact_iter_impl! { u8 i8 + NonZero + // Since RangeInclusive> can only be 1..=uN::MAX the length of this range is always + // <= uN::MAX, so they are always valid ExactSizeIterator unlike the ranges that include zero. + NonZero NonZero // These are incorrect per the reasoning above, // but removing them would be a breaking change as they were stabilized in Rust 1.26.0. diff --git a/library/coretests/tests/iter/range.rs b/library/coretests/tests/iter/range.rs index d5d2b8bf2b047..4a00e6f96bda8 100644 --- a/library/coretests/tests/iter/range.rs +++ b/library/coretests/tests/iter/range.rs @@ -502,3 +502,142 @@ fn test_double_ended_range() { panic!("unreachable"); } } + +macro_rules! nz { + (NonZero<$type:ident>($val:literal)) => { + ::core::num::NonZero::<$type>::new($val).unwrap() + }; + (NonZero<$type:ident>::MAX) => { + ::core::num::NonZero::<$type>::new($type::MAX).unwrap() + }; +} + +macro_rules! nonzero_array { + (NonZero<$type:ident>[$($val:literal),*]) => { + [$(nz!(NonZero<$type>($val))),*] + } +} + +macro_rules! nonzero_range { + (NonZero<$type:ident>($($left:literal)?..$($right:literal)?)) => { + nz!(NonZero<$type>($($left)?))..nz!(NonZero<$type>($($right)?)) + }; + (NonZero<$type:ident>($($left:literal)?..MAX)) => { + nz!(NonZero<$type>($($left)?))..nz!(NonZero<$type>::MAX) + }; + (NonZero<$type:ident>($($left:literal)?..=$right:literal)) => { + nz!(NonZero<$type>($($left)?))..=nz!(NonZero<$type>($right)) + }; + (NonZero<$type:ident>($($left:literal)?..=MAX)) => { + nz!(NonZero<$type>($($left)?))..=nz!(NonZero<$type>::MAX) + }; +} + +#[test] +fn test_nonzero_range() { + assert_eq!( + nonzero_range!(NonZero(1..=21)).step_by(5).collect::>(), + nonzero_array!(NonZero[1, 6, 11, 16, 21]) + ); + assert_eq!( + nonzero_range!(NonZero(1..=20)).step_by(5).collect::>(), + nonzero_array!(NonZero[1, 6, 11, 16]) + ); + assert_eq!( + nonzero_range!(NonZero(1..20)).step_by(5).collect::>(), + nonzero_array!(NonZero[1, 6, 11, 16]) + ); + assert_eq!( + nonzero_range!(NonZero(1..21)).rev().step_by(5).collect::>(), + nonzero_array!(NonZero[20, 15, 10, 5]) + ); + assert_eq!( + nonzero_range!(NonZero(1..21)).rev().step_by(6).collect::>(), + nonzero_array!(NonZero[20, 14, 8, 2]) + ); + assert_eq!( + nonzero_range!(NonZero(200..255)).step_by(50).collect::>(), + nonzero_array!(NonZero[200, 250]) + ); + assert_eq!( + nonzero_range!(NonZero(200..5)).step_by(1).collect::>(), + nonzero_array!(NonZero[]) + ); + assert_eq!( + nonzero_range!(NonZero(200..200)).step_by(1).collect::>(), + nonzero_array!(NonZero[]) + ); + + assert_eq!(nonzero_range!(NonZero(10..20)).step_by(1).size_hint(), (10, Some(10))); + assert_eq!(nonzero_range!(NonZero(10..20)).step_by(5).size_hint(), (2, Some(2))); + assert_eq!(nonzero_range!(NonZero(1..21)).rev().step_by(5).size_hint(), (4, Some(4))); + assert_eq!(nonzero_range!(NonZero(1..21)).rev().step_by(6).size_hint(), (4, Some(4))); + assert_eq!(nonzero_range!(NonZero(20..1)).step_by(1).size_hint(), (0, Some(0))); + assert_eq!(nonzero_range!(NonZero(20..20)).step_by(1).size_hint(), (0, Some(0))); + + // ExactSizeIterator + assert_eq!(nonzero_range!(NonZero(1..=MAX)).step_by(1).len(), usize::from(u8::MAX)); + assert_eq!(nonzero_range!(NonZero(1..=MAX)).step_by(1).len(), usize::from(u16::MAX)); + assert_eq!(nonzero_range!(NonZero(1..=MAX)).step_by(1).len(), usize::MAX); + + // Limits (next) + let mut range = nonzero_range!(NonZero(254..=MAX)); + assert_eq!(range.next(), Some(nz!(NonZero(254)))); + assert_eq!(range.next(), Some(nz!(NonZero(255)))); + assert_eq!(range.next(), None); + + let mut range = nonzero_range!(NonZero(65534..=MAX)); + assert_eq!(range.next(), Some(nz!(NonZero(65534)))); + assert_eq!(range.next(), Some(nz!(NonZero(65535)))); + assert_eq!(range.next(), None); + + // Limits (size_hint, exclusive range) + assert_eq!( + nonzero_range!(NonZero(1..MAX)).step_by(1).size_hint(), + (u8::MAX as usize - 1, Some(u8::MAX as usize - 1)) + ); + assert_eq!( + nonzero_range!(NonZero(1..MAX)).step_by(1).size_hint(), + (u16::MAX as usize - 1, Some(u16::MAX as usize - 1)) + ); + #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] + assert_eq!( + nonzero_range!(NonZero(1..MAX)).step_by(1).size_hint(), + (u32::MAX as usize - 1, Some(u32::MAX as usize - 1)) + ); + #[cfg(target_pointer_width = "64")] + assert_eq!( + nonzero_range!(NonZero(1..MAX)).step_by(1).size_hint(), + (u64::MAX as usize - 1, Some(u64::MAX as usize - 1)) + ); + assert_eq!(nonzero_range!(NonZero(1..MAX)).step_by(1).size_hint(), (usize::MAX, None)); + assert_eq!( + nonzero_range!(NonZero(1..MAX)).step_by(1).size_hint(), + (usize::MAX - 1, Some(usize::MAX - 1)) + ); + + // Limits (size_hint, inclusive range) + assert_eq!( + nonzero_range!(NonZero(1..=MAX)).step_by(1).size_hint(), + (u8::MAX as usize, Some(u8::MAX as usize)) + ); + assert_eq!( + nonzero_range!(NonZero(1..=MAX)).step_by(1).size_hint(), + (u16::MAX as usize, Some(u16::MAX as usize)) + ); + #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] + assert_eq!( + nonzero_range!(NonZero(1..=MAX)).step_by(1).size_hint(), + (u32::MAX as usize, Some(u32::MAX as usize)) + ); + #[cfg(target_pointer_width = "64")] + assert_eq!( + nonzero_range!(NonZero(1..=MAX)).step_by(1).size_hint(), + (u64::MAX as usize, Some(u64::MAX as usize)) + ); + assert_eq!(nonzero_range!(NonZero(1..=MAX)).step_by(1).size_hint(), (usize::MAX, None)); + assert_eq!( + nonzero_range!(NonZero(1..=MAX)).step_by(1).size_hint(), + (usize::MAX, Some(usize::MAX)) + ); +}