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,