diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 625b67a79ad..f5756da6a30 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -305,10 +305,13 @@ pub fn reserve_for_push(&mut self, len: usize) { /// The same as `reserve`, but returns on errors instead of panicking or aborting. pub fn try_reserve(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> { if self.needs_to_grow(len, additional) { - self.grow_amortized(len, additional) - } else { - Ok(()) + self.grow_amortized(len, additional)?; } + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + core::intrinsics::assume(!self.needs_to_grow(len, additional)); + } + Ok(()) } /// Ensures that the buffer contains at least enough space to hold `len + @@ -339,7 +342,14 @@ pub fn try_reserve_exact( len: usize, additional: usize, ) -> Result<(), TryReserveError> { - if self.needs_to_grow(len, additional) { self.grow_exact(len, additional) } else { Ok(()) } + if self.needs_to_grow(len, additional) { + self.grow_exact(len, additional)?; + } + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + core::intrinsics::assume(!self.needs_to_grow(len, additional)); + } + Ok(()) } /// Shrinks the buffer down to the specified capacity. If the given amount diff --git a/tests/codegen/vec-reserve-extend.rs b/tests/codegen/vec-reserve-extend.rs new file mode 100644 index 00000000000..d95220104c2 --- /dev/null +++ b/tests/codegen/vec-reserve-extend.rs @@ -0,0 +1,14 @@ +// compile-flags: -O + +#![crate_type = "lib"] + +// CHECK-LABEL: @should_reserve_once +#[no_mangle] +pub fn should_reserve_once(v: &mut Vec) { + // CHECK: tail call void @llvm.assume + v.try_reserve(3).unwrap(); + // CHECK-NOT: call {{.*}}reserve + // CHECK-NOT: call {{.*}}do_reserve_and_handle + // CHECK-NOT: call {{.*}}__rust_alloc( + v.extend([1, 2, 3]); +} diff --git a/tests/ui/hygiene/panic-location.run.stderr b/tests/ui/hygiene/panic-location.run.stderr index 1fd61084130..e0dc13c0c95 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:535:5: +thread 'main' panicked at library/alloc/src/raw_vec.rs:545:5: capacity overflow note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace