Optimize `str.escape_default().to_string()`

@@ -1,6 +1,5 @@ //! Code related to parsing literals. -use std::fmt::Write as _; use std::{ascii, fmt, str};  use rustc_literal_escaper::{@@ -38,19 +37,8 @@ pub fn escape_string_symbol(symbol: Symbol) -> Symbol {     let mut escaped = String::with_capacity(s.len() + 1);     escaped.push_str(prefix); -    // Don't use escape_default() here, because using it is slower than escaping manually.-    for c in suffix.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(),-        }-    }+    escaped.extend(suffix.escape_default());+     Symbol::intern(&escaped) } 

@@ -92,6 +92,7 @@ #![feature(allocator_ext)] #![feature(array_into_iter_constructors)] #![feature(ascii_char)]+#![feature(ascii_char_variants)] #![feature(async_fn_traits)] #![feature(async_iterator)] #![feature(borrowed_buf_init)]@@ -156,6 +157,7 @@ #![feature(pattern)] #![feature(pin_coerce_unsized_trait)] #![feature(ptr_alignment_type)]+#![feature(ptr_cast_array)] #![feature(ptr_metadata)] #![feature(random)] #![feature(raw_os_error_ty)]

@@ -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;@@ -2448,12 +2452,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 {@@ -2549,10 +2576,7 @@ impl<'a> FromIterator<&'a core::ascii::Char> for String { #[stable(feature = "rust1", since = "1.0.0")] impl Extend<char> for String {     fn extend<I: IntoIterator<Item = char>>(&mut self, iter: I) {-        let iterator = iter.into_iter();-        let (lower_bound, _) = iterator.size_hint();-        self.reserve(lower_bound);-        iterator.for_each(move |c| self.push(c));+        SpecExtendChars::spec_extend_into(iter, self);     }      #[inline]@@ -2566,6 +2590,28 @@ impl Extend<char> for String {     } } +#[cfg(not(no_global_oom_handling))]+trait SpecExtendChars {+    fn spec_extend_into(self, s: &mut String);+}++#[cfg(not(no_global_oom_handling))]+impl<T: IntoIterator<Item = char>> SpecExtendChars for T {+    default fn spec_extend_into(self, target: &mut String) {+        let iterator = self.into_iter();+        let (lower_bound, _) = iterator.size_hint();+        target.reserve(lower_bound);+        iterator.for_each(move |c| target.push(c));+    }+}++#[cfg(not(no_global_oom_handling))]+impl SpecExtendChars for core::str::EscapeDefault<'_> {+    fn spec_extend_into(self, target: &mut String) {+        target.extend_from_escape_default(self);+    }+}+ #[cfg(not(no_global_oom_handling))] #[stable(feature = "extend_ref", since = "1.2.0")] impl<'a> Extend<&'a char> for String {@@ -3160,6 +3206,126 @@ impl SpecToString for fmt::Arguments<'_> {     } } +#[cfg(not(no_global_oom_handling))]+impl SpecToString for core::str::EscapeDefault<'_> {+    fn spec_to_string(&self) -> String {+        let mut s = String::new();+        s.extend_from_escape_default(self.clone());+        s+    }+}++#[cfg(not(no_global_oom_handling))]+impl String {+    #[inline]+    fn extend_from_escape_default(&mut self, escaper: core::str::EscapeDefault<'_>) {+        let (front, middle, tail) = escaper.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);+        self.reserve(lengths);++        if let Some(front) = front {+            self.extend(front);+        }+        if let Some(middle) = middle {+            /// Ensures a spare capacity of N bytes and then passes that as MaybeUninit+            /// slice to a closure that returns the bytes written, then+            /// updates the length.+            ///+            /// # 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();+                // SAFETY: We reserved enough space for the requested number of bytes.+                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);+                }+            }++            for c in middle {+                match c {+                    '\'' | '\"' | '\\' => {+                        self.push('\\');+                        self.push(c);+                    }+                    '\x20'..='\x7e' => self.push(c),+                    '\t' => {+                        self.push('\\');+                        self.push('t');+                    }+                    '\r' => {+                        self.push('\\');+                        self.push('r');+                    }+                    '\n' => {+                        self.push('\\');+                        self.push('n');+                    }+                    c => {+                        // SAFETY: We only let the function write ascii chars, so it remains valid utf8.+                        // and it returns the number of bytes written, so so with_spare ensures the rest for us.+                        unsafe {+                            with_spare::<10, ascii::Char>(self.as_mut_vec(), |spare| {+                                escape_unicode_into(spare, c)+                            });+                        }+                    }+                }+            }+        }++        if let Some(tail) = tail {+            self.extend(tail);+        }+    }+}++#[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]

@@ -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";

@@ -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())+}

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

@@ -1589,6 +1589,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)]