From 029fbd67ef65dae50172678033c0762610041a38 Mon Sep 17 00:00:00 2001
From: Kornel <kornel@geekhood.net>
Date: Mon, 9 Oct 2023 15:05:32 +0100
Subject: [PATCH] Hint optimizer about reserved capacity

---
 library/alloc/src/raw_vec.rs               | 18 ++++++++++++++----
 tests/codegen/vec-reserve-extend.rs        | 14 ++++++++++++++
 tests/ui/hygiene/panic-location.run.stderr |  2 +-
 3 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 tests/codegen/vec-reserve-extend.rs

diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs
index 625b67a79ad09..f5756da6a309e 100644
--- a/library/alloc/src/raw_vec.rs
+++ b/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
diff --git a/tests/codegen/vec-reserve-extend.rs b/tests/codegen/vec-reserve-extend.rs
new file mode 100644
index 0000000000000..d95220104c229
--- /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<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]);
+}
diff --git a/tests/ui/hygiene/panic-location.run.stderr b/tests/ui/hygiene/panic-location.run.stderr
index 1fd61084130a9..e0dc13c0c95c8 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