Rework reserve assumes for Vec and String

library/alloc/src/raw_vec.rs

@@ -308,6 +308,10 @@ impl<T, A: Allocator> RawVec<T, A> {
         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<T, A: Allocator> RawVec<T, A> {
         }
         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<T, A: Allocator> RawVec<T, A> {
         }
         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<T, A: Allocator> RawVec<T, A> {
 impl<T, A: Allocator> RawVec<T, A> {
     /// 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)
     }
 

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.

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<T, A: Allocator> Vec<T, A> {
     /// ```
     #[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<T, A: Allocator> Vec<T, A> {
     /// ```
     #[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<T, A: Allocator> Vec<T, A> {
     /// # 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<T, A: Allocator> Vec<T, A> {
     /// # 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.

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);
+}

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<u8>, 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<u8>, 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<u8>, 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<u8>, 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);
+}

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