optimize str.replace

formatting

rename some variables

refactor replace_ascii

fix function variable names

Author: LaihoE

library/alloc/src/str.rs

@@ -268,7 +268,33 @@ impl str {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[inline]
     pub fn replace<P: Pattern>(&self, from: P, to: &str) -> String {
-        let mut result = String::new();
+        let mut capacity = 0;
+
+        #[inline]
+        fn get_minimum_result_capacity(input_len: usize, from_len: usize, to_len: usize) -> usize {
+            // check if output is going to be at least as long as input.
+            if from_len <= to_len {
+                return input_len;
+            }
+            // lower bound where we have the maximum number of matches:
+            // max_n_matches = len(input) / len(from)
+            // capacity = max_n_matches * to_utf8_len
+            let max_n_matches = input_len.checked_div(from_len).unwrap_or(0);
+            return max_n_matches * to_len;
+        }
+
+        // Path for patterns that can be represented as utf8 bytes (str, char etc.).
+        if let Some(from_as_utf8) = from.as_utf8_bytes() {
+            let from_utf8_len = from_as_utf8.len();
+            let to_utf8_len = to.as_bytes().len();
+            // Fast path for ascii
+            if from_utf8_len == 1 && to_utf8_len == 1 {
+                return replace_ascii(&self.as_bytes(), from_as_utf8[0], to.as_bytes()[0]);
+            }
+            capacity = get_minimum_result_capacity(self.bytes().len(), from_utf8_len, to_utf8_len);
+        }
+
+        let mut result = String::with_capacity(capacity);
         let mut last_end = 0;
         for (start, part) in self.match_indices(from) {
             result.push_str(unsafe { self.get_unchecked(last_end..start) });
@@ -661,3 +687,11 @@ fn convert_while_ascii(b: &[u8], convert: fn(&u8) -> u8) -> Vec<u8> {
 
     out
 }
+#[inline]
+#[cfg(not(test))]
+#[cfg(not(no_global_oom_handling))]
+fn replace_ascii(utf8_bytes: &[u8], from: u8, to: u8) -> String {
+    let result: Vec<u8> = utf8_bytes.iter().map(|b| if *b == from { to } else { *b }).collect();
+    // SAFETY: We replaced ascii with ascii on valid utf8 strings.
+    unsafe { String::from_utf8_unchecked(result) }
+}

library/alloc/src/string.rs

@@ -2319,6 +2319,11 @@ impl<'b> Pattern for &'b String {
     {
         self[..].strip_suffix_of(haystack)
     }
+
+    #[inline]
+    fn as_utf8_bytes(&self) -> Option<&[u8]> {
+        Some(self.as_bytes())
+    }
 }
 
 macro_rules! impl_eq {

library/core/src/str/pattern.rs

@@ -160,6 +160,9 @@ pub trait Pattern: Sized {
             None
         }
     }
+
+    /// Returns the pattern as utf-8 bytes if possible.
+    fn as_utf8_bytes(&self) -> Option<&[u8]>;
 }
 
 // Searcher
@@ -599,6 +602,11 @@ impl Pattern for char {
     {
         self.encode_utf8(&mut [0u8; 4]).strip_suffix_of(haystack)
     }
+
+    #[inline]
+    fn as_utf8_bytes(&self) -> Option<&[u8]> {
+        None
+    }
 }
 
 /////////////////////////////////////////////////////////////////////////////
@@ -657,6 +665,11 @@ impl<C: MultiCharEq> Pattern for MultiCharEqPattern<C> {
     fn into_searcher(self, haystack: &str) -> MultiCharEqSearcher<'_, C> {
         MultiCharEqSearcher { haystack, char_eq: self.0, char_indices: haystack.char_indices() }
     }
+
+    #[inline]
+    fn as_utf8_bytes(&self) -> Option<&[u8]> {
+        None
+    }
 }
 
 unsafe impl<'a, C: MultiCharEq> Searcher<'a> for MultiCharEqSearcher<'a, C> {
@@ -747,6 +760,11 @@ macro_rules! pattern_methods {
         {
             ($pmap)(self).strip_suffix_of(haystack)
         }
+
+        #[inline]
+        fn as_utf8_bytes(&self) -> Option<&[u8]> {
+            None
+        }
     };
 }
 
@@ -1022,6 +1040,11 @@ impl<'b> Pattern for &'b str {
             None
         }
     }
+
+    #[inline]
+    fn as_utf8_bytes(&self) -> Option<&[u8]> {
+        Some(self.as_bytes())
+    }
 }
 
 /////////////////////////////////////////////////////////////////////////////