optimize Cstr/EscapeAscii display

old:
    ascii::bench_ascii_escape_display_mixed      17.97µs/iter +/- 204.00ns
    ascii::bench_ascii_escape_display_no_escape 545.00ns/iter   +/- 6.00ns
new:
    ascii::bench_ascii_escape_display_mixed      4.99µs/iter +/- 56.00ns
    ascii::bench_ascii_escape_display_no_escape 91.00ns/iter  +/- 1.00ns

5 files changed, 67 insertions, 1 deletions

diff --git a/library/core/src/ascii.rs b/library/core/src/ascii.rs
index ef8e4d098ed..02867789b79 100644
--- a/library/core/src/ascii.rs
+++ b/library/core/src/ascii.rs
@@ -96,6 +96,17 @@ pub fn escape_default(c: u8) -> EscapeDefault {
     EscapeDefault(escape::EscapeIterInner::new(data, range))
 }
 
+impl EscapeDefault {
+    pub(crate) fn empty() -> Self {
+        let data = [Char::Null; 4];
+        EscapeDefault(escape::EscapeIterInner::new(data, 0..0))
+    }
+
+    pub(crate) fn as_str(&self) -> &str {
+        self.0.as_str()
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl Iterator for EscapeDefault {
     type Item = u8;
diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs
index 2568aaf34f3..f3992b500ad 100644
--- a/library/core/src/iter/adapters/flatten.rs
+++ b/library/core/src/iter/adapters/flatten.rs
@@ -18,6 +18,14 @@ impl<I: Iterator, U: IntoIterator, F: FnMut(I::Item) -> U> FlatMap<I, U, F> {
     pub(in crate::iter) fn new(iter: I, f: F) -> FlatMap<I, U, F> {
         FlatMap { inner: FlattenCompat::new(iter.map(f)) }
     }
+
+    pub(crate) fn into_parts(self) -> (Option<U::IntoIter>, Option<I>, Option<U::IntoIter>) {
+        (
+            self.inner.frontiter,
+            self.inner.iter.into_inner().map(Map::into_inner),
+            self.inner.backiter,
+        )
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/library/core/src/iter/adapters/fuse.rs b/library/core/src/iter/adapters/fuse.rs
index b1fa4f92117..e38234eae49 100644
--- a/library/core/src/iter/adapters/fuse.rs
+++ b/library/core/src/iter/adapters/fuse.rs
@@ -24,6 +24,10 @@ impl<I> Fuse<I> {
     pub(in crate::iter) fn new(iter: I) -> Fuse<I> {
         Fuse { iter: Some(iter) }
     }
+
+    pub(crate) fn into_inner(self) -> Option<I> {
+        self.iter
+    }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
diff --git a/library/core/src/iter/adapters/map.rs b/library/core/src/iter/adapters/map.rs
index 31d02a4da6e..2563f27d169 100644
--- a/library/core/src/iter/adapters/map.rs
+++ b/library/core/src/iter/adapters/map.rs
@@ -68,6 +68,10 @@ impl<I, F> Map<I, F> {
     pub(in crate::iter) fn new(iter: I, f: F) -> Map<I, F> {
         Map { iter, f }
     }
+
+    pub(crate) fn into_inner(self) -> I {
+        self.iter
+    }
 }
 
 #[stable(feature = "core_impl_debug", since = "1.9.0")]
diff --git a/library/core/src/slice/ascii.rs b/library/core/src/slice/ascii.rs
index f3311f76a7f..324bf5c05af 100644
--- a/library/core/src/slice/ascii.rs
+++ b/library/core/src/slice/ascii.rs
@@ -5,6 +5,7 @@ use crate::fmt::{self, Write};
 use crate::iter;
 use crate::mem;
 use crate::ops;
+use core::ascii::EscapeDefault;
 
 #[cfg(not(test))]
 impl [u8] {
@@ -250,7 +251,45 @@ impl<'a> iter::FusedIterator for EscapeAscii<'a> {}
 #[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
 impl<'a> fmt::Display for EscapeAscii<'a> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.clone().try_for_each(|b| f.write_char(b as char))
+        // disassemble iterator, including front/back parts of flatmap in case it has been partially consumed
+        let (front, slice, back) = self.clone().inner.into_parts();
+        let front = front.unwrap_or(EscapeDefault::empty());
+        let mut bytes = slice.unwrap_or_default().as_slice();
+        let back = back.unwrap_or(EscapeDefault::empty());
+
+        // usually empty, so the formatter won't have to do any work
+        for byte in front {
+            f.write_char(byte as char)?;
+        }
+
+        fn needs_escape(b: u8) -> bool {
+            b > 0x7E || b < 0x20 || b == b'\\' || b == b'\'' || b == b'"'
+        }
+
+        while bytes.len() > 0 {
+            // fast path for the printable, non-escaped subset of ascii
+            let prefix = bytes.iter().take_while(|&&b| !needs_escape(b)).count();
+            // SAFETY: prefix length was derived by counting bytes in the same splice, so it's in-bounds
+            let (prefix, remainder) = unsafe { bytes.split_at_unchecked(prefix) };
+            // SAFETY: prefix is a valid utf8 sequence, as it's a subset of ASCII
+            let prefix = unsafe { crate::str::from_utf8_unchecked(prefix) };
+
+            f.write_str(prefix)?; // the fast part
+
+            bytes = remainder;
+
+            if let Some(&b) = bytes.first() {
+                // guaranteed to be non-empty, better to write it as a str
+                f.write_str(ascii::escape_default(b).as_str())?;
+                bytes = &bytes[1..];
+            }
+        }
+
+        // also usually empty
+        for byte in back {
+            f.write_char(byte as char)?;
+        }
+        Ok(())
     }
 }
 #[stable(feature = "inherent_ascii_escape", since = "1.60.0")]