Skip to content

Commit

Permalink
Add StrExt::replace_smolstr, replacen_smolstr
Browse files Browse the repository at this point in the history
  • Loading branch information
alexheretic authored and Veykril committed Apr 18, 2024
1 parent 9971a3f commit 1e2145d
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 5 deletions.
47 changes: 42 additions & 5 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -592,6 +592,22 @@ pub trait StrExt: private::Sealed {
/// See [`str::to_ascii_uppercase`].
#[must_use = "this returns a new SmolStr without modifying the original"]
fn to_ascii_uppercase_smolstr(&self) -> SmolStr;

/// Replaces all matches of a &str with another &str returning a new [`SmolStr`],
/// potentially without allocating.
///
/// See [`str::replace`].
// TODO: Use `Pattern` when stable.
#[must_use = "this returns a new SmolStr without modifying the original"]
fn replace_smolstr(&self, from: &str, to: &str) -> SmolStr;

/// Replaces first N matches of a &str with another &str returning a new [`SmolStr`],
/// potentially without allocating.
///
/// See [`str::replacen`].
// TODO: Use `Pattern` when stable.
#[must_use = "this returns a new SmolStr without modifying the original"]
fn replacen_smolstr(&self, from: &str, to: &str, count: usize) -> SmolStr;
}

impl StrExt for str {
Expand All @@ -614,6 +630,24 @@ impl StrExt for str {
fn to_ascii_uppercase_smolstr(&self) -> SmolStr {
SmolStr::from_char_iter(self.chars().map(|c| c.to_ascii_uppercase()))
}

#[inline]
fn replace_smolstr(&self, from: &str, to: &str) -> SmolStr {
self.replacen_smolstr(from, to, usize::MAX)
}

#[inline]
fn replacen_smolstr(&self, from: &str, to: &str, count: usize) -> SmolStr {
let mut result = Writer::new();
let mut last_end = 0;
for (start, part) in self.match_indices(from).take(count) {
result.push_str(unsafe { self.get_unchecked(last_end..start) });
result.push_str(to);
last_end = start + part.len();
}
result.push_str(unsafe { self.get_unchecked(last_end..self.len()) });
SmolStr::from(result)
}
}

mod private {
Expand Down Expand Up @@ -651,10 +685,8 @@ impl Writer {
len: 0,
}
}
}

impl fmt::Write for Writer {
fn write_str(&mut self, s: &str) -> fmt::Result {
fn push_str(&mut self, s: &str) {
// if currently on the stack
if self.len <= INLINE_CAP {
let old_len = self.len;
Expand All @@ -663,8 +695,7 @@ impl fmt::Write for Writer {
// if the new length will fit on the stack (even if it fills it entirely)
if self.len <= INLINE_CAP {
self.inline[old_len..self.len].copy_from_slice(s.as_bytes());

return Ok(()); // skip the heap push below
return; // skip the heap push below
}

self.heap.reserve(self.len);
Expand All @@ -678,7 +709,13 @@ impl fmt::Write for Writer {
}

self.heap.push_str(s);
}
}

impl fmt::Write for Writer {
#[inline]
fn write_str(&mut self, s: &str) -> fmt::Result {
self.push_str(s);
Ok(())
}
}
Expand Down
14 changes: 14 additions & 0 deletions tests/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -312,4 +312,18 @@ mod test_str_ext {
assert_eq!(uppercase, "AßΔC");
assert!(!uppercase.is_heap_allocated());
}

#[test]
fn replace() {
let result = "foo_bar_baz".replace_smolstr("ba", "do");
assert_eq!(result, "foo_dor_doz");
assert!(!result.is_heap_allocated());
}

#[test]
fn replacen() {
let result = "foo_bar_baz".replacen_smolstr("ba", "do", 1);
assert_eq!(result, "foo_dor_baz");
assert!(!result.is_heap_allocated());
}
}

0 comments on commit 1e2145d

Please sign in to comment.