refactor RangeFromIter overflow-checks impl

library/core/src/range/iter.rs

@@ -298,9 +298,9 @@ range_incl_exact_iter_impl! {
 #[derive(Debug, Clone)]
 pub struct RangeFromIter<A> {
     start: A,
-    /// Whether the first element of the iterator has yielded.
+    /// Whether the maximum value of the iterator has yielded.
     /// Only used when overflow checks are enabled.
-    first: bool,
+    exhausted: bool,
 }
 
 impl<A: Step> RangeFromIter<A> {
@@ -309,10 +309,12 @@ impl<A: Step> RangeFromIter<A> {
     #[rustc_inherit_overflow_checks]
     #[unstable(feature = "new_range_api", issue = "125687")]
     pub fn remainder(self) -> RangeFrom<A> {
-        if intrinsics::overflow_checks() {
-            if !self.first {
-                return RangeFrom { start: Step::forward(self.start, 1) };
-            }
+        // Need to handle this case even if overflow-checks are disabled,
+        // because a `RangeFromIter` could be exhausted in a crate with
+        // overflow-checks enabled, but then passed to a crate with them
+        // disabled before this is called.
+        if self.exhausted {
+            return RangeFrom { start: Step::forward(self.start, 1) };
         }
 
         RangeFrom { start: self.start }
@@ -326,14 +328,29 @@ impl<A: Step> Iterator for RangeFromIter<A> {
     #[inline]
     #[rustc_inherit_overflow_checks]
     fn next(&mut self) -> Option<A> {
+        if self.exhausted {
+            // This should panic if overflow checks are enabled, since
+            // `forward_checked` returned `None` in prior iteration.
+            self.start = Step::forward(self.start.clone(), 1);
+
+            // If we get here, if means this iterator was exhausted by a crate
+            // with overflow-checks enabled, but now we're iterating in a crate with
+            // overflow-checks disabled. Since we successfully incremented `self.start`
+            // above (in many cases this will wrap around to MIN), we now unset
+            // the flag so we don't repeat this process in the next iteration.
+            //
+            // This could also happen if `forward_checked` returned None but
+            // (for whatever reason, not applicable to any std implementors)
+            // `forward` doesn't panic when overflow-checks are enabled. In that
+            // case, this is also the correct behavior.
+            self.exhausted = false;
+        }
         if intrinsics::overflow_checks() {
-            if self.first {
-                self.first = false;
+            let Some(n) = Step::forward_checked(self.start.clone(), 1) else {
+                self.exhausted = true;
                 return Some(self.start.clone());
-            }
-
-            self.start = Step::forward(self.start.clone(), 1);
-            return Some(self.start.clone());
+            };
+            return Some(mem::replace(&mut self.start, n));
         }
 
         let n = Step::forward(self.start.clone(), 1);
@@ -348,18 +365,22 @@ impl<A: Step> Iterator for RangeFromIter<A> {
     #[inline]
     #[rustc_inherit_overflow_checks]
     fn nth(&mut self, n: usize) -> Option<A> {
+        // Typically `forward` will cause an overflow-check panic here,
+        // but unset the exhausted flag to handle the uncommon cases.
+        // See the comments in `next` for more details.
+        if self.exhausted {
+            self.start = Step::forward(self.start.clone(), 1);
+            self.exhausted = false;
+        }
         if intrinsics::overflow_checks() {
-            if self.first {
-                self.first = false;
-
-                let plus_n = Step::forward(self.start.clone(), n);
+            let plus_n = Step::forward(self.start.clone(), n);
+            if let Some(plus_n1) = Step::forward_checked(plus_n.clone(), 1) {
+                self.start = plus_n1;
+            } else {
                 self.start = plus_n.clone();
-                return Some(plus_n);
+                self.exhausted = true;
             }
-
-            let plus_n = Step::forward(self.start.clone(), n);
-            self.start = Step::forward(plus_n.clone(), 1);
-            return Some(self.start.clone());
+            return Some(plus_n);
         }
 
         let plus_n = Step::forward(self.start.clone(), n);
@@ -380,6 +401,6 @@ impl<A: Step> IntoIterator for RangeFrom<A> {
     type IntoIter = RangeFromIter<A>;
 
     fn into_iter(self) -> Self::IntoIter {
-        RangeFromIter { start: self.start, first: true }
+        RangeFromIter { start: self.start, exhausted: false }
     }
 }

tests/codegen-llvm/fromrangeiter-overflow-checks.rs

@@ -2,7 +2,7 @@
 // runtime check that panics after yielding the maximum value of the range bound type. That is
 // tested for by tests/ui/iterators/rangefrom-overflow-overflow-checks.rs
 //
-// This test ensures that such a runtime check is *not* emitted when debug-assertions are
+// This test ensures such runtime checks are optimized out when debug-assertions are
 // enabled, but overflow-checks are explicitly disabled.
 
 //@ revisions: DEBUG NOCHECKS
@@ -11,17 +11,24 @@
 
 #![crate_type = "lib"]
 #![feature(new_range_api)]
-use std::range::{RangeFrom, RangeFromIter};
+use std::range::RangeFrom;
 
-// CHECK-LABEL: @iterrangefrom_remainder(
+// CHECK-LABEL: @rangefrom_increments(
 #[no_mangle]
-pub unsafe fn iterrangefrom_remainder(x: RangeFromIter<i32>) -> RangeFrom<i32> {
-    //        DEBUG: i32 noundef %x
-    //     NOCHECKS: i32 noundef returned %x
-    //        DEBUG: br i1
-    //        DEBUG: call core::panicking::panic_const::panic_const_add_overflow
-    //        DEBUG: unreachable
-    // NOCHECKS-NOT: unreachable
-    //     NOCHECKS: ret i32 %x
-    x.remainder()
+pub unsafe fn rangefrom_increments(range: RangeFrom<i32>) -> RangeFrom<i32> {
+    // Iterator is contained entirely within this function, so the optimizer should
+    // be able to see that `exhausted` is never set and optimize out any branches.
+
+    //         CHECK: i32 noundef %range
+    //         DEBUG: switch i32 %range
+    //         DEBUG: call core::panicking::panic_const::panic_const_add_overflow
+    //         DEBUG: unreachable
+    //  NOCHECKS-NOT: unreachable
+    //      NOCHECKS: [[REM:%[a-z_0-9.]+]] = add i32 %range, 2
+    // NOCHECKS-NEXT: ret i32 [[REM]]
+
+    let mut iter = range.into_iter();
+    let _ = iter.next();
+    let _ = iter.next();
+    iter.remainder()
 }

tests/ui/iterators/auxiliary/rangefrom-overflow-2crates-ocno.rs

@@ -0,0 +1,10 @@
+//@ compile-flags: -C overflow-checks=no
+
+#![crate_type = "lib"]
+#![feature(new_range_api)]
+
+use std::range::RangeFromIter;
+
+pub fn next(iter: &mut RangeFromIter<u8>) -> u8 {
+    iter.next().unwrap()
+}

tests/ui/iterators/auxiliary/rangefrom-overflow-2crates-ocyes.rs

@@ -0,0 +1,10 @@
+//@ compile-flags: -C overflow-checks=yes
+
+#![crate_type = "lib"]
+#![feature(new_range_api)]
+
+use std::range::RangeFromIter;
+
+pub fn next(iter: &mut RangeFromIter<u8>) -> u8 {
+    iter.next().unwrap()
+}

tests/ui/iterators/rangefrom-overflow-2crates.rs

@@ -0,0 +1,40 @@
+//@ run-pass
+//@ needs-unwind
+//@ aux-build:rangefrom-overflow-2crates-ocno.rs
+//@ aux-build:rangefrom-overflow-2crates-ocyes.rs
+
+// For #154124
+// Test that two crates with different overflow-checks have the same results,
+// even when the iterator is passed between them.
+
+#![feature(new_range_api)]
+
+extern crate rangefrom_overflow_2crates_ocno;
+extern crate rangefrom_overflow_2crates_ocyes;
+
+use rangefrom_overflow_2crates_ocno::next as next_ocno;
+use rangefrom_overflow_2crates_ocyes::next as next_ocyes;
+
+fn main() {
+    let mut iter_ocyes = std::range::RangeFrom::from(0_u8..).into_iter();
+    let mut iter_ocno = iter_ocyes.clone();
+
+    for n in 0_u8..=255 {
+        assert_eq!(n, next_ocno(&mut iter_ocyes.clone()));
+        assert_eq!(n, next_ocyes(&mut iter_ocyes));
+        assert_eq!(n, next_ocyes(&mut iter_ocno.clone()));
+        assert_eq!(n, next_ocno(&mut iter_ocno));
+    }
+
+    // `iter_ocno` should have wrapped
+    assert_eq!(0, next_ocyes(&mut iter_ocno.clone()));
+    assert_eq!(0, next_ocno(&mut iter_ocno));
+    // `iter_ocyes` should be exhausted,
+    // which will wrap when called without overflow-checks
+    assert_eq!(0, next_ocno(&mut iter_ocyes.clone()));
+    // and panic when called with overflow-checks
+    let r = std::panic::catch_unwind(move || {
+        let _ = next_ocyes(&mut iter_ocyes);
+    });
+    assert!(r.is_err());
+}