Rollup merge of #92124 - AngelicosPhosphoros:remove_extra_alloc_in_cstring_new_35838, r=Mark-Simulacrum

Little improves in CString `new` when creating from slice

Old code already contain optimization for cases with `&str` and `&[u8]` args. This commit adds a specialization for `&mut[u8]` too.

Also, I added usage of old slice in search for zero bytes instead of new buffer because it produce better code for constant inputs on Windows LTO builds. For other platforms, this wouldn't cause any difference because it calls `libc` anyway.

Inlined `_new` method into spec trait to reduce amount of code generated to `CString::new` callers.

1 files changed, 45 insertions, 22 deletions

diff --git a/library/std/src/ffi/c_str.rs b/library/std/src/ffi/c_str.rs
index 9c1b79d6966..d859bff1a45 100644
--- a/library/std/src/ffi/c_str.rs
+++ b/library/std/src/ffi/c_str.rs
@@ -373,38 +373,61 @@ impl CString {
     /// the position of the nul byte.
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn new<T: Into<Vec<u8>>>(t: T) -> Result<CString, NulError> {
-        trait SpecIntoVec {
-            fn into_vec(self) -> Vec<u8>;
+        trait SpecNewImpl {
+            fn spec_new_impl(self) -> Result<CString, NulError>;
         }
-        impl<T: Into<Vec<u8>>> SpecIntoVec for T {
-            default fn into_vec(self) -> Vec<u8> {
-                self.into()
+
+        impl<T: Into<Vec<u8>>> SpecNewImpl for T {
+            default fn spec_new_impl(self) -> Result<CString, NulError> {
+                let bytes: Vec<u8> = self.into();
+                match memchr::memchr(0, &bytes) {
+                    Some(i) => Err(NulError(i, bytes)),
+                    None => Ok(unsafe { CString::from_vec_unchecked(bytes) }),
+                }
             }
         }
-        // Specialization for avoiding reallocation.
-        impl SpecIntoVec for &'_ [u8] {
-            fn into_vec(self) -> Vec<u8> {
-                let mut v = Vec::with_capacity(self.len() + 1);
-                v.extend(self);
-                v
+
+        // Specialization for avoiding reallocation
+        #[inline(always)] // Without that it is not inlined into specializations
+        fn spec_new_impl_bytes(bytes: &[u8]) -> Result<CString, NulError> {
+            // We cannot have such large slice that we would overflow here
+            // but using `checked_add` allows LLVM to assume that capacity never overflows
+            // and generate twice shorter code.
+            // `saturating_add` doesn't help for some reason.
+            let capacity = bytes.len().checked_add(1).unwrap();
+
+            // Allocate before validation to avoid duplication of allocation code.
+            // We still need to allocate and copy memory even if we get an error.
+            let mut buffer = Vec::with_capacity(capacity);
+            buffer.extend(bytes);
+
+            // Check memory of self instead of new buffer.
+            // This allows better optimizations if lto enabled.
+            match memchr::memchr(0, bytes) {
+                Some(i) => Err(NulError(i, buffer)),
+                None => Ok(unsafe { CString::from_vec_unchecked(buffer) }),
             }
         }
-        impl SpecIntoVec for &'_ str {
-            fn into_vec(self) -> Vec<u8> {
-                let mut v = Vec::with_capacity(self.len() + 1);
-                v.extend(self.as_bytes());
-                v
+
+        impl SpecNewImpl for &'_ [u8] {
+            fn spec_new_impl(self) -> Result<CString, NulError> {
+                spec_new_impl_bytes(self)
             }
         }
 
-        Self::_new(SpecIntoVec::into_vec(t))
-    }
+        impl SpecNewImpl for &'_ str {
+            fn spec_new_impl(self) -> Result<CString, NulError> {
+                spec_new_impl_bytes(self.as_bytes())
+            }
+        }
 
-    fn _new(bytes: Vec<u8>) -> Result<CString, NulError> {
-        match memchr::memchr(0, &bytes) {
-            Some(i) => Err(NulError(i, bytes)),
-            None => Ok(unsafe { CString::from_vec_unchecked(bytes) }),
+        impl SpecNewImpl for &'_ mut [u8] {
+            fn spec_new_impl(self) -> Result<CString, NulError> {
+                spec_new_impl_bytes(self)
+            }
         }
+
+        t.spec_new_impl()
     }
 
     /// Creates a C-compatible string by consuming a byte vector,