Optimise StrExt::replace_smolstr, StrExt::replacen_smolstr for single ascii replace

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## Unreleased
+
+- Optimise `StrExt::replace_smolstr`, `StrExt::replacen_smolstr` for single ascii replace,
+  ~3x speedup inline & heap.
+
 ## 0.3.2 - 2024-10-23
 
 - Fix `SmolStrBuilder::push` incorrectly padding null bytes when spilling onto the heap on a

src/lib.rs

@@ -658,7 +658,28 @@ impl StrExt for str {
     }
 
     #[inline]
-    fn replacen_smolstr(&self, from: &str, to: &str, count: usize) -> SmolStr {
+    fn replacen_smolstr(&self, from: &str, to: &str, mut count: usize) -> SmolStr {
+        // Fast path for replacing a single ASCII character with another inline.
+        if let [from_u8] = from.as_bytes() {
+            if let [to_u8] = to.as_bytes() {
+                return if self.len() <= count {
+                    // SAFETY: `from_u8` & `to_u8` are ascii
+                    unsafe { replacen_1_ascii(self, |b| if b == from_u8 { *to_u8 } else { *b }) }
+                } else {
+                    unsafe {
+                        replacen_1_ascii(self, |b| {
+                            if b == from_u8 && count != 0 {
+                                count -= 1;
+                                *to_u8
+                            } else {
+                                *b
+                            }
+                        })
+                    }
+                };
+            }
+        }
+
         let mut result = SmolStrBuilder::new();
         let mut last_end = 0;
         for (start, part) in self.match_indices(from).take(count) {
@@ -675,6 +696,26 @@ impl StrExt for str {
     }
 }
 
+/// SAFETY: `map` fn must only replace ascii with ascii or return unchanged bytes.
+#[inline]
+unsafe fn replacen_1_ascii(src: &str, mut map: impl FnMut(&u8) -> u8) -> SmolStr {
+    if src.len() <= INLINE_CAP {
+        let mut buf = [0u8; INLINE_CAP];
+        for (idx, b) in src.as_bytes().iter().enumerate() {
+            buf[idx] = map(b);
+        }
+        SmolStr(Repr::Inline {
+            // SAFETY: `len` is in bounds
+            len: unsafe { InlineSize::transmute_from_u8(src.len() as u8) },
+            buf,
+        })
+    } else {
+        let out = src.as_bytes().iter().map(map).collect();
+        // SAFETY: We replaced ascii with ascii on valid utf8 strings.
+        unsafe { String::from_utf8_unchecked(out).into() }
+    }
+}
+
 impl<T> ToSmolStr for T
 where
     T: fmt::Display + ?Sized,

tests/test.rs

@@ -389,6 +389,13 @@ mod test_str_ext {
         assert_eq!(result, "foo_dor_baz");
         assert!(!result.is_heap_allocated());
     }
+
+    #[test]
+    fn replacen_1_ascii() {
+        let result = "foo_bar_baz".replacen_smolstr("o", "u", 1);
+        assert_eq!(result, "fuo_bar_baz");
+        assert!(!result.is_heap_allocated());
+    }
 }
 
 #[cfg(feature = "borsh")]