new RangeFrom overflow check

Author: pitaj

library/core/src/range/iter.rs

@@ -3,6 +3,7 @@ use crate::iter::{
 };
 use crate::num::NonZero;
 use crate::range::{Range, RangeFrom, RangeInclusive, legacy};
+use crate::{intrinsics, mem};
 
 /// By-value [`Range`] iterator.
 #[unstable(feature = "new_range_api", issue = "125687")]
@@ -293,12 +294,25 @@ range_incl_exact_iter_impl! {
 /// By-value [`RangeFrom`] iterator.
 #[unstable(feature = "new_range_api", issue = "125687")]
 #[derive(Debug, Clone)]
-pub struct IterRangeFrom<A>(legacy::RangeFrom<A>);
+pub struct IterRangeFrom<A> {
+    start: A,
+    /// Whether the first element of the iterator has yielded.
+    /// Only used when overflow checks are enabled.
+    first: bool,
+}
 
-impl<A> IterRangeFrom<A> {
+impl<A: Step> IterRangeFrom<A> {
     /// Returns the remainder of the range being iterated over.
+    #[inline]
+    #[rustc_inherit_overflow_checks]
     pub fn remainder(self) -> RangeFrom<A> {
-        RangeFrom { start: self.0.start }
+        if intrinsics::overflow_checks() {
+            if !self.first {
+                return RangeFrom { start: Step::forward(self.start, 1) };
+            }
+        }
+
+        RangeFrom { start: self.start }
     }
 }
 
@@ -307,18 +321,47 @@ impl<A: Step> Iterator for IterRangeFrom<A> {
     type Item = A;
 
     #[inline]
+    #[rustc_inherit_overflow_checks]
     fn next(&mut self) -> Option<A> {
-        self.0.next()
+        if intrinsics::overflow_checks() {
+            if self.first {
+                self.first = false;
+                return Some(self.start.clone());
+            }
+
+            self.start = Step::forward(self.start.clone(), 1);
+            return Some(self.start.clone());
+        }
+
+        let n = Step::forward(self.start.clone(), 1);
+        Some(mem::replace(&mut self.start, n))
     }
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        self.0.size_hint()
+        (usize::MAX, None)
     }
 
     #[inline]
+    #[rustc_inherit_overflow_checks]
     fn nth(&mut self, n: usize) -> Option<A> {
-        self.0.nth(n)
+        if intrinsics::overflow_checks() {
+            if self.first {
+                self.first = false;
+
+                let plus_n = Step::forward(self.start.clone(), n);
+                self.start = plus_n.clone();
+                return Some(plus_n);
+            }
+
+            let plus_n = Step::forward(self.start.clone(), n);
+            self.start = Step::forward(plus_n.clone(), 1);
+            return Some(self.start.clone());
+        }
+
+        let plus_n = Step::forward(self.start.clone(), n);
+        self.start = Step::forward(plus_n.clone(), 1);
+        Some(plus_n)
     }
 }
 
@@ -334,6 +377,6 @@ impl<A: Step> IntoIterator for RangeFrom<A> {
     type IntoIter = IterRangeFrom<A>;
 
     fn into_iter(self) -> Self::IntoIter {
-        IterRangeFrom(self.into())
+        IterRangeFrom { start: self.start, first: true }
     }
 }

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

@@ -0,0 +1,27 @@
+// With -Coverflow-checks=yes (enabled by default by -Cdebug-assertions=yes) we will produce a
+// 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
+// enabled, but overflow-checks are explicitly disabled.
+
+//@ revisions: DEBUG NOCHECKS
+//@ compile-flags: -O -Cdebug-assertions=yes
+//@ [NOCHECKS] compile-flags: -Coverflow-checks=no
+
+#![crate_type = "lib"]
+#![feature(new_range_api)]
+use std::range::{IterRangeFrom, RangeFrom};
+
+// CHECK-LABEL: @iterrangefrom_remainder(
+#[no_mangle]
+pub unsafe fn iterrangefrom_remainder(x: IterRangeFrom<i32>) -> RangeFrom<i32> {
+    //    DEBUG: i32 noundef %x.0
+    // NOCHECKS: i32 noundef returned %x.0
+    //    CHECK: i1 noundef zeroext %x.1
+    //    DEBUG: br i1
+    //    DEBUG: call core::panicking::panic_const::panic_const_add_overflow
+    //    DEBUG: unreachable
+    // NOCHECKS: ret i32 %x.0
+    x.remainder()
+}

tests/ui/iterators/iterrangefrom.rs

@@ -0,0 +1,23 @@
+//@ run-pass
+//@ compile-flags: -C overflow-checks=yes
+
+#![feature(new_range_api)]
+
+use std::{iter, range};
+
+fn main() {
+    for (a, b) in iter::zip(0_u32..256, range::RangeFrom::from(0_u8..)) {
+        assert_eq!(a, u32::from(b));
+    }
+
+    let mut a = range::RangeFrom::from(0_u8..).into_iter();
+    let mut b = 0_u8..;
+    assert_eq!(a.next(), b.next());
+    assert_eq!(a.nth(5), b.nth(5));
+    assert_eq!(a.nth(0), b.next());
+
+    let mut a = range::RangeFrom::from(0_u8..).into_iter();
+    let mut b = 0_u8..;
+    assert_eq!(a.nth(5), b.nth(5));
+    assert_eq!(a.nth(0), b.next());
+}

tests/ui/iterators/rangefrom-overflow-debug.rs

@@ -0,0 +1,17 @@
+//@ run-pass
+//@ needs-unwind
+//@ compile-flags: -O -C debug_assertions=yes
+
+#![feature(new_range_api)]
+
+use std::panic;
+
+fn main() {
+    let mut it = core::range::RangeFrom::from(u8::MAX..).into_iter();
+    assert_eq!(it.next().unwrap(), 255);
+
+    let r = panic::catch_unwind(|| {
+        let _ = it.remainder();
+    });
+    assert!(r.is_err());
+}

tests/ui/iterators/rangefrom-overflow-ndebug.rs

@@ -0,0 +1,10 @@
+//@ run-pass
+//@ compile-flags: -O -C debug_assertions=no
+
+#![feature(new_range_api)]
+
+fn main() {
+    let mut it = core::range::RangeFrom::from(u8::MAX..).into_iter();
+    assert_eq!(it.next().unwrap(), 255);
+    assert_eq!(it.remainder().start, u8::MIN);
+}

tests/ui/iterators/rangefrom-overflow-overflow-checks.rs

@@ -0,0 +1,17 @@
+//@ run-pass
+//@ needs-unwind
+//@ compile-flags: -O -C overflow-checks=yes
+
+#![feature(new_range_api)]
+
+use std::panic;
+
+fn main() {
+    let mut it = core::range::RangeFrom::from(u8::MAX..).into_iter();
+    assert_eq!(it.next().unwrap(), 255);
+
+    let r = panic::catch_unwind(|| {
+        let _ = it.remainder();
+    });
+    assert!(r.is_err());
+}