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.

Author: jalil-salame

library/core/src/iter/range.rs

@@ -16,6 +16,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.
 ///
@@ -431,6 +432,138 @@ step_integer_impls! {
     wider than 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)
+        }
+    };
+}
+
+macro_rules! step_nonzero_impls {
+    {
+        narrower than or same width as usize:
+            $( $narrower:ident ),+;
+        wider than 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 steps_between(start: &Self, end: &Self) -> Option<usize> {
+                    if *start <= *end {
+                        // This relies on $u_narrower <= usize
+                        Some((end.get() - start.get()) as usize)
+                    } else {
+                        None
+                    }
+                }
+
+                #[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 steps_between(start: &Self, end: &Self) -> Option<usize> {
+                    if *start <= *end {
+                        usize::try_from(end.get() - start.get()).ok()
+                    } else {
+                        None
+                    }
+                }
+
+                #[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! {
+    narrower than or same width as usize: u8, u16, u32, u64, usize;
+    wider than usize: u128;
+}
+
+#[cfg(target_pointer_width = "32")]
+step_nonzero_impls! {
+    narrower than or same width as usize: u8, u16, u32, usize;
+    wider than usize: u64, u128;
+}
+
+#[cfg(target_pointer_width = "16")]
+step_nonzero_impls! {
+    narrower than or same width as usize: u8, u16, usize;
+    wider than usize: u32, u64, u128;
+}
+
 #[unstable(feature = "step_trait", reason = "recently redesigned", issue = "42168")]
 impl Step for char {
     #[inline]

library/core/tests/iter/range.rs

@@ -501,3 +501,70 @@ fn test_double_ended_range() {
         panic!("unreachable");
     }
 }
+
+macro_rules! nz {
+    ($type:ident($val:literal)) => {
+        $type::new($val).unwrap()
+    };
+}
+
+macro_rules! nonzero_array {
+    ($type:ident[$($val:literal),*]) => {
+        [$(nz!($type($val))),*]
+    }
+}
+
+macro_rules! nonzero_range {
+    ($type:ident($($left:literal)?..$($right:literal)?)) => {
+        nz!($type($($left)?))..nz!($type($($right)?))
+    };
+    ($type:ident($($left:literal)?..=$right:literal)) => {
+        nz!($type($($left)?))..=nz!($type($right))
+    };
+}
+
+#[test]
+fn test_nonzero_range() {
+    #![allow(deprecated)]
+    use core::num::{NonZeroU32, NonZeroU8, NonZeroUsize};
+
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(1..=21)).step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[1, 6, 11, 16, 21])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(1..=20)).step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[1, 6, 11, 16])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(1..20)).step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[1, 6, 11, 16])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(1..21)).rev().step_by(5).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[20, 15, 10, 5])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(1..21)).rev().step_by(6).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[20, 14, 8, 2])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroU8(200..255)).step_by(50).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroU8[200, 250])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(200..5)).step_by(1).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[])
+    );
+    assert_eq!(
+        nonzero_range!(NonZeroUsize(200..200)).step_by(1).collect::<Vec<_>>(),
+        nonzero_array!(NonZeroUsize[])
+    );
+
+    assert_eq!(nonzero_range!(NonZeroU32(10..20)).step_by(1).size_hint(), (10, Some(10)));
+    assert_eq!(nonzero_range!(NonZeroU32(10..20)).step_by(5).size_hint(), (2, Some(2)));
+    assert_eq!(nonzero_range!(NonZeroU32(1..21)).rev().step_by(5).size_hint(), (4, Some(4)));
+    assert_eq!(nonzero_range!(NonZeroU32(1..21)).rev().step_by(6).size_hint(), (4, Some(4)));
+    assert_eq!(nonzero_range!(NonZeroU32(20..1)).step_by(1).size_hint(), (0, Some(0)));
+    assert_eq!(nonzero_range!(NonZeroU32(20..20)).step_by(1).size_hint(), (0, Some(0)));
+}