Expose SmolStrBuilder

Author: Veykril

src/lib.rs

@@ -620,15 +620,13 @@ pub trait StrExt: private::Sealed {
     /// 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;
 }
@@ -661,7 +659,7 @@ impl StrExt for str {
 
     #[inline]
     fn replacen_smolstr(&self, from: &str, to: &str, count: usize) -> SmolStr {
-        let mut result = Writer::new();
+        let mut result = SmolStrBuilder::new();
         let mut last_end = 0;
         for (start, part) in self.match_indices(from).take(count) {
             // SAFETY: `start` is guaranteed to be within the bounds of `self` as per
@@ -677,6 +675,15 @@ impl StrExt for str {
     }
 }
 
+impl<T> ToSmolStr for T
+where
+    T: fmt::Display + ?Sized,
+{
+    fn to_smolstr(&self) -> SmolStr {
+        format_smolstr!("{}", self)
+    }
+}
+
 mod private {
     /// No downstream impls allowed.
     pub trait Sealed {}
@@ -689,85 +696,94 @@ mod private {
 #[macro_export]
 macro_rules! format_smolstr {
     ($($tt:tt)*) => {{
-        use ::core::fmt::Write;
-        let mut w = $crate::Writer::new();
-        w.write_fmt(format_args!($($tt)*)).expect("a formatting trait implementation returned an error");
-        $crate::SmolStr::from(w)
+        let mut w = $crate::SmolStrBuilder::new();
+        ::core::fmt::Write::write_fmt(&mut w, format_args!($($tt)*)).expect("a formatting trait implementation returned an error");
+        w.finish()
     }};
 }
 
-#[doc(hidden)]
-pub struct Writer {
-    inline: [u8; INLINE_CAP],
-    heap: String,
-    len: usize,
+/// A builder that can be used to efficiently build a [`SmolStr`].
+///
+/// This won't allocate if the final string fits into the inline buffer.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum SmolStrBuilder {
+    Inline { len: usize, buf: [u8; INLINE_CAP] },
+    Heap(String),
+}
+
+impl Default for SmolStrBuilder {
+    #[inline]
+    fn default() -> Self {
+        Self::new()
+    }
 }
 
-impl Writer {
+impl SmolStrBuilder {
+    /// Creates a new empty [`SmolStrBuilder`].
     #[must_use]
     pub const fn new() -> Self {
-        Writer {
-            inline: [0; INLINE_CAP],
-            heap: String::new(),
+        SmolStrBuilder::Inline {
+            buf: [0; INLINE_CAP],
             len: 0,
         }
     }
 
-    fn push_str(&mut self, s: &str) {
-        // if currently on the stack
-        if self.len <= INLINE_CAP {
-            let old_len = self.len;
-            self.len += s.len();
-
-            // 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; // skip the heap push below
+    /// Builds a [`SmolStr`] from `self`.
+    #[must_use]
+    pub fn finish(&self) -> SmolStr {
+        SmolStr(match self {
+            &SmolStrBuilder::Inline { len, buf } => {
+                debug_assert!(len <= INLINE_CAP);
+                Repr::Inline {
+                    // SAFETY: We know that `value.len` is less than or equal to the maximum value of `InlineSize`
+                    len: unsafe { InlineSize::transmute_from_u8(len as u8) },
+                    buf,
+                }
             }
+            SmolStrBuilder::Heap(heap) => Repr::new(heap),
+        })
+    }
 
-            self.heap.reserve(self.len);
-
-            // copy existing inline bytes over to the heap
-            // SAFETY: inline data is guaranteed to be valid utf8 for `old_len` bytes
-            unsafe {
-                self.heap
-                    .as_mut_vec()
-                    .extend_from_slice(&self.inline[..old_len]);
+    /// Appends a given string slice onto the end of `self`'s buffer.
+    pub fn push_str(&mut self, s: &str) {
+        // if currently on the stack
+        match self {
+            Self::Inline { len, buf } => {
+                let old_len = *len;
+                *len += s.len();
+
+                // if the new length will fit on the stack (even if it fills it entirely)
+                if *len <= INLINE_CAP {
+                    buf[old_len..*len].copy_from_slice(s.as_bytes());
+                    return; // skip the heap push below
+                }
+
+                let mut heap = String::with_capacity(*len);
+
+                // copy existing inline bytes over to the heap
+                // SAFETY: inline data is guaranteed to be valid utf8 for `old_len` bytes
+                unsafe {
+                    heap.as_mut_vec().extend_from_slice(&buf[..old_len]);
+                }
+                heap.push_str(s);
+                *self = SmolStrBuilder::Heap(heap);
             }
+            SmolStrBuilder::Heap(heap) => heap.push_str(s),
         }
-
-        self.heap.push_str(s);
     }
 }
 
-impl fmt::Write for Writer {
+impl fmt::Write for SmolStrBuilder {
     #[inline]
     fn write_str(&mut self, s: &str) -> fmt::Result {
         self.push_str(s);
         Ok(())
     }
 }
 
-impl From<Writer> for SmolStr {
-    fn from(value: Writer) -> Self {
-        SmolStr(if value.len <= INLINE_CAP {
-            Repr::Inline {
-                // SAFETY: We know that `value.len` is less than or equal to the maximum value of `InlineSize`
-                len: unsafe { InlineSize::transmute_from_u8(value.len as u8) },
-                buf: value.inline,
-            }
-        } else {
-            Repr::new(&value.heap)
-        })
-    }
-}
-
-impl<T> ToSmolStr for T
-where
-    T: fmt::Display + ?Sized,
-{
-    fn to_smolstr(&self) -> SmolStr {
-        format_smolstr!("{}", self)
+impl From<SmolStrBuilder> for SmolStr {
+    fn from(value: SmolStrBuilder) -> Self {
+        value.finish()
     }
 }
 

tests/test.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 #[cfg(not(miri))]
 use proptest::{prop_assert, prop_assert_eq, proptest};
 
-use smol_str::SmolStr;
+use smol_str::{SmolStr, SmolStrBuilder};
 
 #[test]
 #[cfg(target_pointer_width = "64")]
@@ -255,6 +255,42 @@ fn test_to_smolstr() {
         assert_eq!(a, smol_str::format_smolstr!("{}", a));
     }
 }
+#[test]
+fn test_builder() {
+    //empty
+    let builder = SmolStrBuilder::new();
+    assert_eq!("", builder.finish());
+
+    // inline push
+    let mut builder = SmolStrBuilder::new();
+    builder.push_str("a");
+    builder.push_str("b");
+    let s = builder.finish();
+    assert!(!s.is_heap_allocated());
+    assert_eq!("ab", s);
+
+    // inline max push
+    let mut builder = SmolStrBuilder::new();
+    builder.push_str(&"a".repeat(23));
+    let s = builder.finish();
+    assert!(!s.is_heap_allocated());
+    assert_eq!("a".repeat(23), s);
+
+    // heap push immediate
+    let mut builder = SmolStrBuilder::new();
+    builder.push_str(&"a".repeat(24));
+    let s = builder.finish();
+    assert!(s.is_heap_allocated());
+    assert_eq!("a".repeat(24), s);
+
+    // heap push succession
+    let mut builder = SmolStrBuilder::new();
+    builder.push_str(&"a".repeat(23));
+    builder.push_str(&"a".repeat(23));
+    let s = builder.finish();
+    assert!(s.is_heap_allocated());
+    assert_eq!("a".repeat(46), s);
+}
 
 #[cfg(test)]
 mod test_str_ext {