Hint optimizer about try-reserved capacity

library/alloc/src/raw_vec.rs

@@ -305,10 +305,13 @@ impl<T, A: Allocator> RawVec<T, A> {
  /// 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 @@ impl<T, A: Allocator> RawVec<T, A> {
  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

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<u8>) {
+ // 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]);
+}

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