Auto merge of #149705 - paolobarbolini:string-extend-chunked, r=<try>

Optimize `Extend for String`

Author: rust-bors[bot]

compiler/rustc_ast/src/ast.rs

@@ -181,8 +181,7 @@ pub fn join_path_syms(path: impl IntoIterator<Item = impl Borrow<Symbol>>) -> St
     for sym in iter {
         let sym = *sym.borrow();
         debug_assert_ne!(sym, kw::PathRoot);
-        s.push_str("::");
-        s.push_str(sym.as_str());
+        s.extend(["::", sym.as_str()]);
     }
     s
 }
@@ -201,8 +200,7 @@ pub fn join_path_idents(path: impl IntoIterator<Item = impl Borrow<Ident>>) -> S
     for ident in iter {
         let ident = *ident.borrow();
         debug_assert_ne!(ident.name, kw::PathRoot);
-        s.push_str("::");
-        s.push_str(&ident.to_string());
+        s.extend(["::", &ident.to_string()]);
     }
     s
 }

compiler/rustc_builtin_macros/src/assert/context.rs

@@ -334,9 +334,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
         if self.paths.contains(&path_ident) {
             return;
         } else {
-            self.fmt_string.push_str("  ");
-            self.fmt_string.push_str(path_ident.as_str());
-            self.fmt_string.push_str(" = {:?}\n");
+            self.fmt_string.extend(["  ", path_ident.as_str(), " = {:?}\n"]);
             let _ = self.paths.insert(path_ident);
         }
         let curr_capture_idx = self.capture_decls.len();

library/alloc/src/borrow.rs

@@ -13,8 +13,6 @@ use core::ops::{Deref, DerefPure};
 use Cow::*;
 
 use crate::fmt;
-#[cfg(not(no_global_oom_handling))]
-use crate::string::String;
 
 // FIXME(inference): const bounds removed due to inference regressions found by crater;
 //   see https://github.com/rust-lang/rust/issues/147964
@@ -496,12 +494,14 @@ impl<'a> AddAssign<&'a str> for Cow<'a, str> {
         if self.is_empty() {
             *self = Cow::Borrowed(rhs)
         } else if !rhs.is_empty() {
-            if let Cow::Borrowed(lhs) = *self {
-                let mut s = String::with_capacity(lhs.len() + rhs.len());
-                s.push_str(lhs);
-                *self = Cow::Owned(s);
+            match self {
+                Self::Borrowed(lhs) => {
+                    *self = Cow::Owned([lhs, rhs].into_iter().collect());
+                }
+                Self::Owned(lhs) => {
+                    lhs.push_str(&rhs);
+                }
             }
-            self.to_mut().push_str(rhs);
         }
     }
 }
@@ -513,12 +513,14 @@ impl<'a> AddAssign<Cow<'a, str>> for Cow<'a, str> {
         if self.is_empty() {
             *self = rhs
         } else if !rhs.is_empty() {
-            if let Cow::Borrowed(lhs) = *self {
-                let mut s = String::with_capacity(lhs.len() + rhs.len());
-                s.push_str(lhs);
-                *self = Cow::Owned(s);
+            match self {
+                Self::Borrowed(lhs) => {
+                    *self = Cow::Owned([&*lhs, &*rhs].into_iter().collect());
+                }
+                Self::Owned(lhs) => {
+                    lhs.push_str(&rhs);
+                }
             }
-            self.to_mut().push_str(&rhs);
         }
     }
 }

library/alloc/src/string.rs

@@ -42,6 +42,8 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
+#[cfg(not(no_global_oom_handling))]
+use core::array;
 use core::error::Error;
 use core::iter::FusedIterator;
 #[cfg(not(no_global_oom_handling))]
@@ -2200,6 +2202,36 @@ impl String {
         let slice = self.vec.leak();
         unsafe { from_utf8_unchecked_mut(slice) }
     }
+
+    /// SAFETY: When calling `<S as AsRef<str>>::as_ref()` multiple times, the same value must be returned.
+    #[cfg(not(no_global_oom_handling))]
+    unsafe fn extend_many<S: AsRef<str>>(&mut self, vals: &[S]) {
+        let additional = vals.iter().fold(0usize, |a, s| a.saturating_add(s.as_ref().len()));
+        self.reserve(additional);
+
+        let mut spare = self.vec.spare_capacity_mut().as_mut_ptr().cast_init();
+        for val in vals {
+            let val = val.as_ref();
+            // SAFETY:
+            // - `val` is a valid &str, so `val.as_ptr()` is valid
+            //   for `val.len()` bytes and properly initialized.
+            // - `spare` points to valid spare capacity in the Vec
+            //   with enough space for `val.len()` bytes.
+            //   This is guaranteed because the caller ensures
+            //   that multiple calls to `<S as AsRef<str>>::as_ref()`
+            //   return the same value, and the saturating addition
+            //   stops undercounting by overflow.
+            // - Both pointers are byte-aligned and the regions cannot overlap.
+            unsafe {
+                ptr::copy_nonoverlapping(val.as_ptr(), spare, val.len());
+                spare = spare.add(val.len());
+            }
+        }
+
+        let new_len = self.vec.len() + additional;
+        // SAFETY: the elements have just been initialized
+        unsafe { self.vec.set_len(new_len) }
+    }
 }
 
 impl FromUtf8Error {
@@ -2499,7 +2531,8 @@ impl<'a> Extend<&'a char> for String {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a> Extend<&'a str> for String {
     fn extend<I: IntoIterator<Item = &'a str>>(&mut self, iter: I) {
-        iter.into_iter().for_each(move |s| self.push_str(s));
+        // SAFETY: `<&str as AsRef<str>>::as_ref()` returns the same value when called multiple times
+        unsafe { self.extend_many_chunked(iter.into_iter()) }
     }
 
     #[inline]
@@ -2512,15 +2545,17 @@ impl<'a> Extend<&'a str> for String {
 #[stable(feature = "box_str2", since = "1.45.0")]
 impl<A: Allocator> Extend<Box<str, A>> for String {
     fn extend<I: IntoIterator<Item = Box<str, A>>>(&mut self, iter: I) {
-        iter.into_iter().for_each(move |s| self.push_str(&s));
+        // SAFETY: `<Box<str, A> as AsRef<str>>::as_ref()` returns the same value when called multiple times
+        unsafe { self.extend_many_chunked(iter.into_iter()) }
     }
 }
 
 #[cfg(not(no_global_oom_handling))]
 #[stable(feature = "extend_string", since = "1.4.0")]
 impl Extend<String> for String {
     fn extend<I: IntoIterator<Item = String>>(&mut self, iter: I) {
-        iter.into_iter().for_each(move |s| self.push_str(&s));
+        // SAFETY: `<String as AsRef<str>>::as_ref()` returns the same value when called multiple times
+        unsafe { self.extend_many_chunked(iter.into_iter()) }
     }
 
     #[inline]
@@ -2533,7 +2568,8 @@ impl Extend<String> for String {
 #[stable(feature = "herd_cows", since = "1.19.0")]
 impl<'a> Extend<Cow<'a, str>> for String {
     fn extend<I: IntoIterator<Item = Cow<'a, str>>>(&mut self, iter: I) {
-        iter.into_iter().for_each(move |s| self.push_str(&s));
+        // SAFETY: `<Cow<'a, str> as AsRef<str>>::as_ref()` returns the same value when called multiple times
+        unsafe { self.extend_many_chunked(iter.into_iter()) }
     }
 
     #[inline]
@@ -2570,6 +2606,57 @@ impl<'a> Extend<&'a core::ascii::Char> for String {
     }
 }
 
+#[cfg(not(no_global_oom_handling))]
+trait ExtendManySpec<S, I> {
+    /// SAFETY: When calling `<S as AsRef<str>>::as_ref()` multiple times, the same value must be returned.
+    unsafe fn extend_many_chunked(&mut self, iter: I);
+}
+
+#[cfg(not(no_global_oom_handling))]
+impl<S, I> ExtendManySpec<S, I> for String
+where
+    S: AsRef<str>,
+    I: Iterator<Item = S>,
+{
+    default unsafe fn extend_many_chunked(&mut self, mut iter: I) {
+        let mut repeat = true;
+        while repeat {
+            let chunk = match iter.next_chunk::<8>() {
+                Ok(chunk) => chunk.into_iter(),
+                Err(partial_chunk) => {
+                    repeat = false;
+                    partial_chunk
+                }
+            };
+
+            // SAFETY: the caller ensures that multiple calls to `<S as AsRef<str>>::as_ref()` return the same value.
+            unsafe { self.extend_many(chunk.as_slice()) }
+        }
+    }
+}
+
+#[cfg(not(no_global_oom_handling))]
+impl<S, const N: usize> ExtendManySpec<S, array::IntoIter<S, N>> for String
+where
+    S: AsRef<str>,
+{
+    unsafe fn extend_many_chunked(&mut self, iter: array::IntoIter<S, N>) {
+        // SAFETY: the caller ensures that multiple calls to `<S as AsRef<str>>::as_ref()` return the same value.
+        unsafe { self.extend_many(iter.as_slice()) }
+    }
+}
+
+#[cfg(not(no_global_oom_handling))]
+impl<S> ExtendManySpec<S, vec::IntoIter<S>> for String
+where
+    S: AsRef<str>,
+{
+    unsafe fn extend_many_chunked(&mut self, iter: vec::IntoIter<S>) {
+        // SAFETY: the caller ensures that multiple calls to `<S as AsRef<str>>::as_ref()` return the same value.
+        unsafe { self.extend_many(iter.as_slice()) }
+    }
+}
+
 /// A convenience impl that delegates to the impl for `&str`.
 ///
 /// # Examples

src/librustdoc/html/markdown.rs

@@ -1645,9 +1645,7 @@ pub(crate) fn plain_text_from_events<'a>(
         match &event {
             Event::Text(text) => s.push_str(text),
             Event::Code(code) => {
-                s.push('`');
-                s.push_str(code);
-                s.push('`');
+                s.extend(["`", code, "`"]);
             }
             Event::HardBreak | Event::SoftBreak => s.push(' '),
             Event::Start(Tag::CodeBlock(..)) => break,

src/librustdoc/html/render/print_item.rs

@@ -2257,8 +2257,7 @@ pub(crate) fn compare_names(left: &str, right: &str) -> Ordering {
 
 pub(super) fn full_path(cx: &Context<'_>, item: &clean::Item) -> String {
     let mut s = join_path_syms(&cx.current);
-    s.push_str("::");
-    s.push_str(item.name.unwrap().as_str());
+    s.extend(["::", item.name.unwrap().as_str()]);
     s
 }