Optimize str.escape_default().to_string()

compiler/rustc_ast/src/util/literal.rs

@@ -1,6 +1,5 @@
 //! Code related to parsing literals.
 
-use std::fmt::Write as _;
 use std::{ascii, fmt, str};
 
 use rustc_literal_escaper::{
@@ -15,22 +14,8 @@ use crate::token::{self, Token};
 // Escapes a string, represented as a symbol. Reuses the original symbol,
 // avoiding interning, if no changes are required.
 pub fn escape_string_symbol(symbol: Symbol) -> Symbol {
-    // Don't use escape_default() here, because using it in conjunction with to_string()
-    // is slow.
     let s = symbol.as_str();
-    let mut escaped = String::with_capacity(s.len());
-    for c in s.chars() {
-        match c {
-            '\t' => escaped.push_str("\\t"),
-            '\r' => escaped.push_str("\\r"),
-            '\n' => escaped.push_str("\\n"),
-            '\\' => escaped.push_str("\\\\"),
-            '\'' => escaped.push_str("\\'"),
-            '\"' => escaped.push_str("\\\""),
-            '\x20'..='\x7e' => escaped.push(c),
-            c => write!(escaped, "\\u{{{:x}}}", c as u32).unwrap(),
-        }
-    }
+    let escaped = s.escape_default().to_string();
     if s == escaped { symbol } else { Symbol::intern(&escaped) }
 }
 

library/alloc/src/lib.rs

@@ -92,6 +92,7 @@
 #![feature(allocator_api)]
 #![feature(array_into_iter_constructors)]
 #![feature(ascii_char)]
+#![feature(ascii_char_variants)]
 #![feature(async_fn_traits)]
 #![feature(async_iterator)]
 #![feature(borrowed_buf_init)]
@@ -157,6 +158,7 @@
 #![feature(pattern)]
 #![feature(pin_coerce_unsized_trait)]
 #![feature(ptr_alignment_type)]
+#![feature(ptr_cast_array)]
 #![feature(ptr_cast_slice)]
 #![feature(ptr_internals)]
 #![feature(ptr_metadata)]

library/alloc/src/string.rs

@@ -42,11 +42,15 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
+#[cfg(not(no_global_oom_handling))]
+use core::ascii;
 use core::error::Error;
 use core::iter::FusedIterator;
 #[cfg(not(no_global_oom_handling))]
 use core::iter::from_fn;
 #[cfg(not(no_global_oom_handling))]
+use core::mem::MaybeUninit;
+#[cfg(not(no_global_oom_handling))]
 use core::num::Saturating;
 #[cfg(not(no_global_oom_handling))]
 use core::ops::Add;
@@ -2402,12 +2406,35 @@ impl Clone for String {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl FromIterator<char> for String {
     fn from_iter<I: IntoIterator<Item = char>>(iter: I) -> String {
-        let mut buf = String::new();
-        buf.extend(iter);
+        FromCharIterSpec::to_string_spec(iter.into_iter())
+    }
+}
+
+#[cfg(not(no_global_oom_handling))]
+trait FromCharIterSpec: Iterator<Item = char> {
+    fn to_string_spec(self) -> String;
+}
+
+#[cfg(not(no_global_oom_handling))]
+impl<T> FromCharIterSpec for T
+where
+    T: Iterator<Item = char>,
+{
+    default fn to_string_spec(self) -> String {
+        let mut buf = String::with_capacity(self.size_hint().0);
+        buf.extend(self);
         buf
     }
 }
 
+#[cfg(not(no_global_oom_handling))]
+impl FromCharIterSpec for core::str::EscapeDefault<'_> {
+    fn to_string_spec(self) -> String {
+        // go through the ToString specialization
+        self.to_string()
+    }
+}
+
 #[cfg(not(no_global_oom_handling))]
 #[stable(feature = "string_from_iter_by_ref", since = "1.17.0")]
 impl<'a> FromIterator<&'a char> for String {
@@ -3116,6 +3143,115 @@ impl SpecToString for fmt::Arguments<'_> {
     }
 }
 
+#[cfg(not(no_global_oom_handling))]
+impl SpecToString for core::str::EscapeDefault<'_> {
+    #[inline]
+    fn spec_to_string(&self) -> String {
+        let (front, middle, tail) = self.clone().into_parts();
+        let lengths = front.as_ref().map_or_default(|f| f.size_hint().0)
+            + middle.as_ref().map_or_default(|m| m.size_hint().0)
+            + tail.as_ref().map_or_default(|t| t.size_hint().0);
+        let mut escaped = String::with_capacity(lengths);
+
+        if let Some(front) = front {
+            escaped.extend(front);
+        }
+        if let Some(middle) = middle {
+            unsafe {
+                /// Updates the length of the string in advance and returns the spare capacity of
+                /// N bytes.
+                ///
+                /// # Safety:
+                ///
+                /// The caller must initialize the number of bytes returned by the closure.
+                #[inline]
+                unsafe fn with_spare<const N: usize, T>(
+                    bytes: &mut Vec<u8>,
+                    then: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize,
+                ) {
+                    const {
+                        assert!(core::mem::size_of::<T>() == 1);
+                    }
+                    bytes.reserve(N);
+                    let old_len = bytes.len();
+                    unsafe {
+                        let spare: &mut [MaybeUninit<T>; N] = bytes
+                            .spare_capacity_mut()
+                            .as_mut_ptr()
+                            .cast::<T>()
+                            .cast_uninit()
+                            .cast_array::<N>()
+                            .as_mut_unchecked();
+                        let written = then(spare);
+                        bytes.set_len(old_len + written);
+                    }
+                }
+
+                let escaped = escaped.as_mut_vec();
+                for c in middle {
+                    match c {
+                        '\'' | '\"' | '\\' => {
+                            escaped.push(b'\\');
+                            escaped.push(c as u8);
+                        }
+                        '\x20'..='\x7e' => escaped.push(c as u8),
+                        '\t' => {
+                            escaped.push(b'\\');
+                            escaped.push(b't');
+                        }
+                        '\r' => {
+                            escaped.push(b'\\');
+                            escaped.push(b'r');
+                        }
+                        '\n' => {
+                            escaped.push(b'\\');
+                            escaped.push(b'n');
+                        }
+                        c => {
+                            with_spare::<10, ascii::Char>(escaped, |spare| {
+                                escape_unicode_into(spare, c)
+                            });
+                        }
+                    }
+                }
+            }
+        }
+        if let Some(tail) = tail {
+            escaped.extend(tail);
+        }
+
+        escaped
+    }
+}
+
+#[cfg(not(no_global_oom_handling))]
+#[inline]
+// adapted from core::escape::escape_unicode, but it's left-aligned instead of right-aligned.
+fn escape_unicode_into(output: &mut [MaybeUninit<ascii::Char>; 10], c: char) -> usize {
+    const HEX_DIGITS: [ascii::Char; 16] = *b"0123456789abcdef".as_ascii().unwrap();
+
+    let c = c as u32;
+    // OR-ing `1` ensures that for `c == 0` the code computes that
+    // one digit should be printed.
+    let prefix_skip = (c | 1).leading_zeros() as usize / 4 - 2;
+
+    // Always write all digits, but for shorter representations the unused
+    // digits get overwritten by later writes.
+    output[3usize.saturating_sub(prefix_skip)].write(HEX_DIGITS[((c >> 20) & 15) as usize]);
+    output[4usize.saturating_sub(prefix_skip)].write(HEX_DIGITS[((c >> 16) & 15) as usize]);
+    output[5usize.saturating_sub(prefix_skip)].write(HEX_DIGITS[((c >> 12) & 15) as usize]);
+    output[6usize.saturating_sub(prefix_skip)].write(HEX_DIGITS[((c >> 8) & 15) as usize]);
+    output[7usize.saturating_sub(prefix_skip)].write(HEX_DIGITS[((c >> 4) & 15) as usize]);
+    output[8usize.saturating_sub(prefix_skip)].write(HEX_DIGITS[(c & 15) as usize]);
+
+    output[0].write(ascii::Char::ReverseSolidus);
+    output[1].write(ascii::Char::SmallU);
+    output[2].write(ascii::Char::LeftCurlyBracket);
+    output[9 - prefix_skip].write(ascii::Char::RightCurlyBracket);
+
+    10 - prefix_skip
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl AsRef<str> for String {
     #[inline]

library/alloctests/benches/str.rs

@@ -1,5 +1,9 @@
 use test::{Bencher, black_box};
 
+#[allow(unused)] // this is a symlink to core benches, we don't use everything here
+#[path = "../../coretests/benches/str/corpora.rs"]
+pub(super) mod corpora;
+
 #[bench]
 fn char_iterator(b: &mut Bencher) {
     let s = "ศไทย中华Việt Nam; Mary had a little lamb, Little lamb";

library/alloctests/benches/string.rs

@@ -159,3 +159,15 @@ fn bench_insert_str_long(b: &mut Bencher) {
         x
     })
 }
+
+#[bench]
+fn bench_from_escape_default_ascii(b: &mut Bencher) {
+    let s = super::str::corpora::en::LARGE;
+    b.iter(|| black_box(s).escape_default().to_string())
+}
+
+#[bench]
+fn bench_from_escape_default_multibyte(b: &mut Bencher) {
+    let s = super::str::corpora::emoji::LARGE;
+    b.iter(|| black_box(s).escape_default().to_string())
+}

library/alloctests/tests/string.rs

@@ -956,3 +956,14 @@ fn test_str_concat() {
     let s: String = format!("{a}{b}");
     assert_eq!(s.as_bytes()[9], 'd' as u8);
 }
+
+// alloc has a specialization for EscapeDefault's ToString impl, which isn't covered by core tests.
+#[test]
+fn test_from_escaped_str() {
+    assert_eq!("abcABC012".escape_default().to_string(), "abcABC012");
+    assert_eq!("\t\n\"".escape_default().to_string(), r#"\t\n\""#);
+    assert_eq!("\u{10FFFF}".escape_default().to_string(), "\\u{10ffff}");
+    assert_eq!("\u{0}".escape_default().to_string(), "\\u{0}");
+    assert_eq!("\u{1}".escape_default().to_string(), "\\u{1}");
+    assert_eq!("\u{100}".escape_default().to_string(), "\\u{100}");
+}

library/core/src/str/iter.rs

@@ -1592,6 +1592,18 @@ pub struct EscapeDefault<'a> {
     pub(super) inner: FlatMap<Chars<'a>, char_mod::EscapeDefault, CharEscapeDefault>,
 }
 
+impl<'a> EscapeDefault<'a> {
+    /// Disassembles the iterator into its internal state.
+    /// Used by alloc for an optimized to_string impl.
+    #[unstable(feature = "std_internals", issue = "none")]
+    pub fn into_parts(
+        self,
+    ) -> (Option<impl Iterator<Item = char>>, Option<Chars<'a>>, Option<impl Iterator<Item = char>>)
+    {
+        self.inner.into_parts()
+    }
+}
+
 /// The return type of [`str::escape_unicode`].
 #[stable(feature = "str_escape", since = "1.34.0")]
 #[derive(Clone, Debug)]