Fix an edge case with StepBy::nth on non-fused iterators

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

@@ -254,9 +254,9 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
     default fn spec_nth(&mut self, mut n: usize) -> Option<I::Item> {
         if self.first_take {
             self.first_take = false;
-            let first = self.iter.next();
+            let first = self.iter.next()?;
             if n == 0 {
-                return first;
+                return Some(first);
             }
             n -= 1;
         }
@@ -266,7 +266,7 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
         // n + 1 could overflow
         // thus, if n is usize::MAX, instead of adding one, we call .nth(step)
         if n == usize::MAX {
-            self.iter.nth(step - 1);
+            self.iter.nth(step - 1)?;
         } else {
             n += 1;
         }
@@ -290,7 +290,8 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
                 n -= div_step;
                 nth_step
             };
-            self.iter.nth(nth - 1);
+
+            self.iter.nth(nth - 1)?;
         }
     }
 

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

@@ -93,6 +93,46 @@ fn test_iterator_step_by_nth_overflow() {
     assert_eq!(it.0, (usize::MAX as Bigger) * 1);
 }
 
+#[test]
+#[allow(non_local_definitions)]
+fn test_iterator_step_by_nth_overflow_on_none() {
+    #[cfg(target_pointer_width = "16")]
+    type Bigger = u32;
+    #[cfg(target_pointer_width = "32")]
+    type Bigger = u64;
+    #[cfg(target_pointer_width = "64")]
+    type Bigger = u128;
+
+    #[derive(Clone)]
+    struct Test(Bigger);
+    impl Iterator for &mut Test {
+        type Item = i32;
+        fn next(&mut self) -> Option<Self::Item> {
+            None
+        }
+        fn nth(&mut self, n: usize) -> Option<Self::Item> {
+            self.0 += n as Bigger + 1;
+            None
+        }
+    }
+
+    // usize::MAX * usize::MAX overflow
+    let mut it = Test(0);
+    let mut step_by = (&mut it).step_by(usize::MAX);
+    // first next() call sets StepBy's first_take to false
+    assert_eq!(step_by.next(), None);
+    assert_eq!(step_by.nth(usize::MAX), None);
+    assert_eq!(it.0, usize::MAX as Bigger + 1);
+
+    // usize::MAX * (usize::MAX - 1) overflow
+    let mut it = Test(0);
+    let mut step_by = (&mut it).step_by(usize::MAX);
+    // first next() call sets StepBy's first_take to false
+    assert_eq!(step_by.next(), None);
+    assert_eq!(step_by.nth(usize::MAX - 1), None);
+    assert_eq!(it.0, usize::MAX as Bigger + 1);
+}
+
 #[test]
 fn test_iterator_step_by_nth_try_fold() {
     let mut it = (0..).step_by(10);
@@ -337,3 +377,44 @@ fn test_step_by_new_range_iter() {
     assert_eq!(it.next_back(), Some(10));
     assert_eq!(it.next(), None);
 }
+
+#[test]
+fn test_step_by_nth_non_fused() {
+    struct StepByNthOne {
+        exhausted: bool,
+    }
+    impl Iterator for StepByNthOne {
+        type Item = i32;
+        fn next(&mut self) -> Option<i32> {
+            if self.exhausted {
+                Some(0)
+            } else {
+                self.exhausted = true;
+                None
+            }
+        }
+    }
+
+    let mut iter = StepByNthOne { exhausted: false }.step_by(1);
+    assert_eq!(iter.nth(1), None)
+}
+
+#[test]
+fn test_step_by_nth_non_fused_on_non_first_take() {
+    struct StepByOneNthMax(i32);
+    impl Iterator for StepByOneNthMax {
+        type Item = i32;
+        fn next(&mut self) -> Option<i32> {
+            let prev = self.0;
+            self.0 += 1;
+            if prev == 1 { None } else { Some(prev) }
+        }
+    }
+
+    let mut iter = StepByOneNthMax(0).step_by(1);
+    // sets StepBy iterator first_take field to false
+    assert_eq!(iter.next(), Some(0));
+    // Our underlying iterator should be pointing at a `None` item
+    // so we should expect `StepBy::nth` to return `None`
+    assert_eq!(iter.nth(usize::MAX), None)
+}