diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 74fa30456eb95..a20e9d6d25226 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -308,6 +308,10 @@ impl RawVec { if self.needs_to_grow(len, additional) { do_reserve_and_handle(self, len, additional); } + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(!self.needs_to_grow(len, additional)); + } } /// A specialized version of `reserve()` used only by the hot and @@ -325,7 +329,7 @@ impl RawVec { } unsafe { // Inform the optimizer that the reservation has succeeded or wasn't needed - core::intrinsics::assume(!self.needs_to_grow(len, additional)); + intrinsics::assume(!self.needs_to_grow(len, additional)); } Ok(()) } @@ -363,7 +367,7 @@ impl RawVec { } unsafe { // Inform the optimizer that the reservation has succeeded or wasn't needed - core::intrinsics::assume(!self.needs_to_grow(len, additional)); + intrinsics::assume(!self.needs_to_grow(len, additional)); } Ok(()) } @@ -387,7 +391,7 @@ impl RawVec { impl RawVec { /// Returns if the buffer needs to grow to fulfill the needed extra capacity. /// Mainly used to make inlining reserve-calls possible without inlining `grow`. - fn needs_to_grow(&self, len: usize, additional: usize) -> bool { + pub(crate) fn needs_to_grow(&self, len: usize, additional: usize) -> bool { additional > self.capacity().wrapping_sub(len) } diff --git a/library/alloc/src/string.rs b/library/alloc/src/string.rs index 4d6968157dedd..00d89b929d5b0 100644 --- a/library/alloc/src/string.rs +++ b/library/alloc/src/string.rs @@ -43,7 +43,6 @@ #![stable(feature = "rust1", since = "1.0.0")] use core::error::Error; -use core::fmt; use core::hash; #[cfg(not(no_global_oom_handling))] use core::iter::from_fn; @@ -60,6 +59,7 @@ use core::slice; use core::str::pattern::Pattern; #[cfg(not(no_global_oom_handling))] use core::str::Utf8Chunks; +use core::{fmt, intrinsics}; #[cfg(not(no_global_oom_handling))] use crate::borrow::{Cow, ToOwned}; @@ -1149,7 +1149,11 @@ impl String { #[inline] #[stable(feature = "rust1", since = "1.0.0")] pub fn reserve(&mut self, additional: usize) { - self.vec.reserve(additional) + self.vec.reserve(additional); + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len())); + } } /// Reserves the minimum capacity for at least `additional` bytes more than @@ -1199,7 +1203,11 @@ impl String { #[inline] #[stable(feature = "rust1", since = "1.0.0")] pub fn reserve_exact(&mut self, additional: usize) { - self.vec.reserve_exact(additional) + self.vec.reserve_exact(additional); + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len())); + } } /// Tries to reserve capacity for at least `additional` bytes more than the @@ -1233,8 +1241,14 @@ impl String { /// # process_data("rust").expect("why is the test harness OOMing on 4 bytes?"); /// ``` #[stable(feature = "try_reserve", since = "1.57.0")] + #[inline] pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> { - self.vec.try_reserve(additional) + self.vec.try_reserve(additional)?; + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len())); + } + Ok(()) } /// Tries to reserve the minimum capacity for at least `additional` bytes @@ -1274,8 +1288,14 @@ impl String { /// # process_data("rust").expect("why is the test harness OOMing on 4 bytes?"); /// ``` #[stable(feature = "try_reserve", since = "1.57.0")] + #[inline] pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> { - self.vec.try_reserve_exact(additional) + self.vec.try_reserve_exact(additional)?; + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len())); + } + Ok(()) } /// Shrinks the capacity of this `String` to match its length. diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index e8a096cac869e..562c767425f66 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -56,7 +56,6 @@ #[cfg(not(no_global_oom_handling))] use core::cmp; use core::cmp::Ordering; -use core::fmt; use core::hash::{Hash, Hasher}; #[cfg(not(no_global_oom_handling))] use core::iter; @@ -65,6 +64,7 @@ use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties}; use core::ops::{self, Index, IndexMut, Range, RangeBounds}; use core::ptr::{self, NonNull}; use core::slice::{self, SliceIndex}; +use core::{fmt, intrinsics}; use crate::alloc::{Allocator, Global}; use crate::borrow::{Cow, ToOwned}; @@ -907,8 +907,13 @@ impl Vec { /// ``` #[cfg(not(no_global_oom_handling))] #[stable(feature = "rust1", since = "1.0.0")] + #[inline] pub fn reserve(&mut self, additional: usize) { self.buf.reserve(self.len, additional); + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(!self.buf.needs_to_grow(self.len, additional)); + } } /// Reserves the minimum capacity for at least `additional` more elements to @@ -937,8 +942,13 @@ impl Vec { /// ``` #[cfg(not(no_global_oom_handling))] #[stable(feature = "rust1", since = "1.0.0")] + #[inline] pub fn reserve_exact(&mut self, additional: usize) { self.buf.reserve_exact(self.len, additional); + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(!self.buf.needs_to_grow(self.len, additional)); + } } /// Tries to reserve capacity for at least `additional` more elements to be inserted @@ -974,8 +984,14 @@ impl Vec { /// # process_data(&[1, 2, 3]).expect("why is the test harness OOMing on 12 bytes?"); /// ``` #[stable(feature = "try_reserve", since = "1.57.0")] + #[inline] pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> { - self.buf.try_reserve(self.len, additional) + self.buf.try_reserve(self.len, additional)?; + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(!self.buf.needs_to_grow(self.len, additional)); + } + Ok(()) } /// Tries to reserve the minimum capacity for at least `additional` @@ -1017,8 +1033,14 @@ impl Vec { /// # process_data(&[1, 2, 3]).expect("why is the test harness OOMing on 12 bytes?"); /// ``` #[stable(feature = "try_reserve", since = "1.57.0")] + #[inline] pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> { - self.buf.try_reserve_exact(self.len, additional) + self.buf.try_reserve_exact(self.len, additional)?; + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + intrinsics::assume(!self.buf.needs_to_grow(self.len, additional)); + } + Ok(()) } /// Shrinks the capacity of the vector as much as possible. diff --git a/tests/codegen/string_reserve_push.rs b/tests/codegen/string_reserve_push.rs new file mode 100644 index 0000000000000..ce76e0ed814ec --- /dev/null +++ b/tests/codegen/string_reserve_push.rs @@ -0,0 +1,61 @@ +// compile-flags: -O + +#![crate_type = "lib"] + +use std::arch::asm; + +#[no_mangle] +// CHECK-LABEL: @reserve_push_str( +pub fn reserve_push_str(v: &mut String, s: &str) { + // CHECK: do_reserve_and_handle + // CHECK: "nop" + // CHECK-NOT: do_reserve_and_handle + // CHECK: ret + v.reserve(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.push_str(s); +} + +#[no_mangle] +// CHECK-LABEL: @reserve_exact_push_str( +pub fn reserve_exact_push_str(v: &mut String, s: &str) { + // CHECK: finish_grow + // CHECK: "nop" + // CHECK-NOT: finish_grow + // CHECK: ret + v.reserve_exact(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.push_str(s); +} + +#[no_mangle] +// CHECK-LABEL: @try_reserve_push_str( +pub fn try_reserve_push_str(v: &mut String, s: &str) { + // CHECK: finish_grow + // CHECK: "nop" + // CHECK-NOT: finish_grow + // CHECK: ret + v.try_reserve(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.push_str(s); +} + +#[no_mangle] +// CHECK-LABEL: @try_reserve_exact_push_str( +pub fn try_reserve_exact_push_str(v: &mut String, s: &str) { + // CHECK: finish_grow + // CHECK: "nop" + // CHECK-NOT: finish_grow + // CHECK: ret + v.try_reserve_exact(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.push_str(s); +} diff --git a/tests/codegen/vec_reserve_extend.rs b/tests/codegen/vec_reserve_extend.rs new file mode 100644 index 0000000000000..abd2674a9ca35 --- /dev/null +++ b/tests/codegen/vec_reserve_extend.rs @@ -0,0 +1,61 @@ +// compile-flags: -O + +#![crate_type = "lib"] + +use std::arch::asm; + +#[no_mangle] +// CHECK-LABEL: @reserve_extend( +pub fn reserve_extend(v: &mut Vec, s: &[u8]) { + // CHECK: do_reserve_and_handle + // CHECK: "nop" + // CHECK-NOT: do_reserve_and_handle + // CHECK: ret + v.reserve(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.extend_from_slice(s); +} + +#[no_mangle] +// CHECK-LABEL: @reserve_exact_extend( +pub fn reserve_exact_extend(v: &mut Vec, s: &[u8]) { + // CHECK: finish_grow + // CHECK: "nop" + // CHECK-NOT: finish_grow + // CHECK: ret + v.reserve_exact(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.extend_from_slice(s); +} + +#[no_mangle] +// CHECK-LABEL: @try_reserve_extend( +pub fn try_reserve_extend(v: &mut Vec, s: &[u8]) { + // CHECK: finish_grow + // CHECK: "nop" + // CHECK-NOT: finish_grow + // CHECK: ret + v.try_reserve(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.extend_from_slice(s); +} + +#[no_mangle] +// CHECK-LABEL: @try_reserve_exact_extend( +pub fn try_reserve_exact_extend(v: &mut Vec, s: &[u8]) { + // CHECK: finish_grow + // CHECK: "nop" + // CHECK-NOT: finish_grow + // CHECK: ret + v.try_reserve_exact(s.len()); + unsafe { + asm!("nop"); // Used to make the check logic easier + } + v.extend_from_slice(s); +} diff --git a/tests/ui/hygiene/panic-location.run.stderr b/tests/ui/hygiene/panic-location.run.stderr index 5c552411da7f3..915274fd9ef8d 100644 --- a/tests/ui/hygiene/panic-location.run.stderr +++ b/tests/ui/hygiene/panic-location.run.stderr @@ -1,3 +1,3 @@ -thread 'main' panicked at library/alloc/src/raw_vec.rs:571:5: +thread 'main' panicked at library/alloc/src/raw_vec.rs:575:5: capacity overflow note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace