diff --git a/library/core/src/range/iter.rs b/library/core/src/range/iter.rs index c1d4fbbd23ad4..23ace6e1ba253 100644 --- a/library/core/src/range/iter.rs +++ b/library/core/src/range/iter.rs @@ -298,9 +298,9 @@ range_incl_exact_iter_impl! { #[derive(Debug, Clone)] pub struct RangeFromIter { 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 RangeFromIter { @@ -309,10 +309,12 @@ impl RangeFromIter { #[rustc_inherit_overflow_checks] #[unstable(feature = "new_range_api", issue = "125687")] pub fn remainder(self) -> RangeFrom { - 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 Iterator for RangeFromIter { #[inline] #[rustc_inherit_overflow_checks] fn next(&mut self) -> Option { + 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 Iterator for RangeFromIter { #[inline] #[rustc_inherit_overflow_checks] fn nth(&mut self, n: usize) -> Option { + // 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 IntoIterator for RangeFrom { type IntoIter = RangeFromIter; fn into_iter(self) -> Self::IntoIter { - RangeFromIter { start: self.start, first: true } + RangeFromIter { start: self.start, exhausted: false } } } diff --git a/tests/codegen-llvm/fromrangeiter-overflow-checks.rs b/tests/codegen-llvm/fromrangeiter-overflow-checks.rs index 4d27f118ddd37..455c81c633407 100644 --- a/tests/codegen-llvm/fromrangeiter-overflow-checks.rs +++ b/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) -> RangeFrom { - // 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) -> RangeFrom { + // 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() } diff --git a/tests/ui/iterators/auxiliary/rangefrom-overflow-2crates-ocno.rs b/tests/ui/iterators/auxiliary/rangefrom-overflow-2crates-ocno.rs new file mode 100644 index 0000000000000..ebba4ae92c47f --- /dev/null +++ b/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 { + iter.next().unwrap() +} diff --git a/tests/ui/iterators/auxiliary/rangefrom-overflow-2crates-ocyes.rs b/tests/ui/iterators/auxiliary/rangefrom-overflow-2crates-ocyes.rs new file mode 100644 index 0000000000000..8eb5392c7bac7 --- /dev/null +++ b/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 { + iter.next().unwrap() +} diff --git a/tests/ui/iterators/rangefrom-overflow-2crates.rs b/tests/ui/iterators/rangefrom-overflow-2crates.rs new file mode 100644 index 0000000000000..c35c96f99322e --- /dev/null +++ b/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()); +}