Rollup merge of rust-lang#127534 - jalil-salame:nonzero-uint-step, r=programmerjake,tgross35

feat(core): impl Step for NonZero<u*>

Implement Step for NonZero unsigned integers as discussed in [libs-team#130][1].

[1]: rust-lang/libs-team#130

`step_nonzero_impls` was adapted from `step_integer_impls` and the tests were adapted from the step tests.

It seems to compile and pass the tests c:

I would like to test the edge cases, specially around overflows. To ensure safe code cannot generate a 0 `NonZero`, but I don't know how to go about it. I would love some feedback on that (and the PR overall).

The impls are tagges with `#[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")]`. I assumed this was appropriate, but I am not sure.

Author: jhpratt

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<u8> NonZero<u16> NonZero<u32> NonZero<u64> NonZero<u128> NonZero<usize>];
 
 /// 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<usize>) {
+            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<Self> {
+                    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<Self> {
+                    match $narrower::try_from(n) {
+                        // *_sub() is not implemented on NonZero<T>
+                        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<Self> {
+                    start.checked_add(n as $wider)
+                }
+
+                #[inline]
+                fn backward_checked(start: Self, n: usize) -> Option<Self> {
+                    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<A: Step> Iterator for ops::Range<A> {
 range_exact_iter_impl! {
     usize u8 u16
     isize i8 i16
+    NonZero<usize> NonZero<u8> NonZero<u16>
 
     // 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<usize> NonZero<u8> NonZero<u16>
 }
 
 #[cfg(target_pointer_width = "32")]
 unsafe_range_trusted_random_access_impl! {
     u32 i32
+    NonZero<u32>
 }
 
 #[cfg(target_pointer_width = "64")]
 unsafe_range_trusted_random_access_impl! {
     u32 i32
     u64 i64
+    NonZero<u32>
+    NonZero<u64>
 }
 
 range_incl_exact_iter_impl! {
     u8
     i8
+    NonZero<u8>
+    // Since RangeInclusive<NonZero<uN>> 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<u16> NonZero<usize>
 
     // These are incorrect per the reasoning above,
     // but removing them would be a breaking change as they were stabilized in Rust 1.26.0.

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<usize>(1..=21)).step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[1, 6, 11, 16, 21])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(1..=20)).step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[1, 6, 11, 16])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(1..20)).step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[1, 6, 11, 16])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(1..21)).rev().step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[20, 15, 10, 5])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(1..21)).rev().step_by(6).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[20, 14, 8, 2])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<u8>(200..255)).step_by(50).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<u8>[200, 250])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(200..5)).step_by(1).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[])
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(200..200)).step_by(1).collect::<Vec<_>>(),
+        nonzero_array!(NonZero<usize>[])
+    );
+
+    assert_eq!(nonzero_range!(NonZero<u32>(10..20)).step_by(1).size_hint(), (10, Some(10)));
+    assert_eq!(nonzero_range!(NonZero<u32>(10..20)).step_by(5).size_hint(), (2, Some(2)));
+    assert_eq!(nonzero_range!(NonZero<u32>(1..21)).rev().step_by(5).size_hint(), (4, Some(4)));
+    assert_eq!(nonzero_range!(NonZero<u32>(1..21)).rev().step_by(6).size_hint(), (4, Some(4)));
+    assert_eq!(nonzero_range!(NonZero<u32>(20..1)).step_by(1).size_hint(), (0, Some(0)));
+    assert_eq!(nonzero_range!(NonZero<u32>(20..20)).step_by(1).size_hint(), (0, Some(0)));
+
+    // ExactSizeIterator
+    assert_eq!(nonzero_range!(NonZero<u8>(1..=MAX)).step_by(1).len(), usize::from(u8::MAX));
+    assert_eq!(nonzero_range!(NonZero<u16>(1..=MAX)).step_by(1).len(), usize::from(u16::MAX));
+    assert_eq!(nonzero_range!(NonZero<usize>(1..=MAX)).step_by(1).len(), usize::MAX);
+
+    // Limits (next)
+    let mut range = nonzero_range!(NonZero<u8>(254..=MAX));
+    assert_eq!(range.next(), Some(nz!(NonZero<u8>(254))));
+    assert_eq!(range.next(), Some(nz!(NonZero<u8>(255))));
+    assert_eq!(range.next(), None);
+
+    let mut range = nonzero_range!(NonZero<u16>(65534..=MAX));
+    assert_eq!(range.next(), Some(nz!(NonZero<u16>(65534))));
+    assert_eq!(range.next(), Some(nz!(NonZero<u16>(65535))));
+    assert_eq!(range.next(), None);
+
+    // Limits (size_hint, exclusive range)
+    assert_eq!(
+        nonzero_range!(NonZero<u8>(1..MAX)).step_by(1).size_hint(),
+        (u8::MAX as usize - 1, Some(u8::MAX as usize - 1))
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<u16>(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<u32>(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<u64>(1..MAX)).step_by(1).size_hint(),
+        (u64::MAX as usize - 1, Some(u64::MAX as usize - 1))
+    );
+    assert_eq!(nonzero_range!(NonZero<u128>(1..MAX)).step_by(1).size_hint(), (usize::MAX, None));
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(1..MAX)).step_by(1).size_hint(),
+        (usize::MAX - 1, Some(usize::MAX - 1))
+    );
+
+    // Limits (size_hint, inclusive range)
+    assert_eq!(
+        nonzero_range!(NonZero<u8>(1..=MAX)).step_by(1).size_hint(),
+        (u8::MAX as usize, Some(u8::MAX as usize))
+    );
+    assert_eq!(
+        nonzero_range!(NonZero<u16>(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<u32>(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<u64>(1..=MAX)).step_by(1).size_hint(),
+        (u64::MAX as usize, Some(u64::MAX as usize))
+    );
+    assert_eq!(nonzero_range!(NonZero<u128>(1..=MAX)).step_by(1).size_hint(), (usize::MAX, None));
+    assert_eq!(
+        nonzero_range!(NonZero<usize>(1..=MAX)).step_by(1).size_hint(),
+        (usize::MAX, Some(usize::MAX))
+    );
+}