Add extend_one, extend_reserve and ensure_capacity

Cargo.toml

@@ -27,6 +27,7 @@ harness = false
 [features]
 default = ["std"]
 std = []
+unstable = []
 test = ["std", "arbitrary", "arbitrary/derive"]
 
 [dependencies]

src/boxed.rs

@@ -32,11 +32,16 @@ pub(crate) struct BoxedString {
 ///
 /// Returns `true` if aligned to an odd address, `false` if even. The sense of
 /// the boolean is "does this look like an InlineString? true/false"
+#[inline]
 fn check_alignment(ptr: *const u8) -> bool {
     ptr.align_offset(2) > 0
 }
 
 impl GenericString for BoxedString {
+    fn cap(&self) -> usize {
+        self.cap
+    }
+
     fn set_size(&mut self, size: usize) {
         self.len = size;
         debug_assert!(self.len <= self.cap);
@@ -53,6 +58,7 @@ impl GenericString for BoxedString {
 impl BoxedString {
     const MINIMAL_CAPACITY: usize = MAX_INLINE * 2;
 
+    #[inline]
     pub(crate) fn check_alignment(this: &Self) -> bool {
         check_alignment(this.ptr.as_ptr())
     }

src/inline.rs

@@ -56,6 +56,10 @@ impl DerefMut for InlineString {
 }
 
 impl GenericString for InlineString {
+    fn cap(&self) -> usize {
+        MAX_INLINE
+    }
+
     fn set_size(&mut self, size: usize) {
         self.marker.set_data(size as u8);
     }

src/lib.rs

@@ -88,6 +88,7 @@
 //! | Feature | Description |
 //! | ------- | ----------- |
 //! | [`arbitrary`](https://crates.io/crates/arbitrary) | [`Arbitrary`][Arbitrary] implementation for [`SmartString`]. |
+//! | `unstable` | Features only available with a nightly Rust compiler. Currently this is only the [`extend_reserve`][core::iter::traits::collect::Extend::extend_reserve] implementation for [`SmartString`]. |
 //! | [`proptest`](https://crates.io/crates/proptest) | A strategy for generating [`SmartString`]s from a regular expression. |
 //! | [`serde`](https://crates.io/crates/serde) | [`Serialize`][Serialize] and [`Deserialize`][Deserialize] implementations for [`SmartString`]. |
 //!
@@ -102,6 +103,7 @@
 #![warn(unreachable_pub, missing_debug_implementations, missing_docs)]
 #![cfg_attr(not(feature = "std"), no_std)]
 #![cfg_attr(needs_allocator_feature, feature(allocator_api))]
+#![cfg_attr(feature = "unstable", feature(extend_one))]
 
 extern crate alloc;
 
@@ -421,6 +423,17 @@ impl<Mode: SmartStringMode> SmartString<Mode> {
         }
     }
 
+    /// Ensures that the string has a capacity of at least the given number of bytes.
+    /// This function will reallocate the string (and therefore unbox a boxed string)
+    /// in order to fit the new capacity.
+    ///
+    /// Note that if the string's capacity is already large enough, this function does nothing.
+    /// This also applies to inline strings, when `capacity` is less than or equal to
+    /// [`MAX_INLINE`].
+    pub fn ensure_capacity(&mut self, target_cap: usize) {
+        string_op_grow!(ops::EnsureCapacity, self, target_cap)
+    }
+
     /// Push a character to the end of the string.
     pub fn push(&mut self, ch: char) {
         string_op_grow!(ops::Push, self, ch)
@@ -710,6 +723,11 @@ impl<'a, Mode: SmartStringMode> Extend<&'a str> for SmartString<Mode> {
             self.push_str(item);
         }
     }
+
+    #[cfg(feature = "unstable")]
+    fn extend_one(&mut self, item: &'a str) {
+        self.push_str(item)
+    }
 }
 
 impl<'a, Mode: SmartStringMode> Extend<&'a char> for SmartString<Mode> {
@@ -718,6 +736,16 @@ impl<'a, Mode: SmartStringMode> Extend<&'a char> for SmartString<Mode> {
             self.push(*item);
         }
     }
+
+    #[cfg(feature = "unstable")]
+    fn extend_one(&mut self, item: &'a char) {
+        self.push(*item)
+    }
+
+    #[cfg(feature = "unstable")]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.ensure_capacity(self.capacity() + additional)
+    }
 }
 
 impl<Mode: SmartStringMode> Extend<char> for SmartString<Mode> {
@@ -726,6 +754,16 @@ impl<Mode: SmartStringMode> Extend<char> for SmartString<Mode> {
             self.push(item);
         }
     }
+
+    #[cfg(feature = "unstable")]
+    fn extend_one(&mut self, item: char) {
+        self.push(item)
+    }
+
+    #[cfg(feature = "unstable")]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.ensure_capacity(self.capacity() + additional)
+    }
 }
 
 impl<Mode: SmartStringMode> Extend<SmartString<Mode>> for SmartString<Mode> {
@@ -742,6 +780,11 @@ impl<Mode: SmartStringMode> Extend<String> for SmartString<Mode> {
             self.push_str(&item);
         }
     }
+
+    #[cfg(feature = "unstable")]
+    fn extend_one(&mut self, item: String) {
+        self.push_str(&item)
+    }
 }
 
 impl<'a, Mode: SmartStringMode + 'a> Extend<&'a SmartString<Mode>> for SmartString<Mode> {
@@ -750,6 +793,11 @@ impl<'a, Mode: SmartStringMode + 'a> Extend<&'a SmartString<Mode>> for SmartStri
             self.push_str(item);
         }
     }
+
+    #[cfg(feature = "unstable")]
+    fn extend_one(&mut self, item: &'a SmartString<Mode>) {
+        self.push_str(item)
+    }
 }
 
 impl<'a, Mode: SmartStringMode> Extend<&'a String> for SmartString<Mode> {
@@ -758,6 +806,11 @@ impl<'a, Mode: SmartStringMode> Extend<&'a String> for SmartString<Mode> {
             self.push_str(item);
         }
     }
+
+    #[cfg(feature = "unstable")]
+    fn extend_one(&mut self, item: &'a String) {
+        self.push_str(item)
+    }
 }
 
 impl<Mode: SmartStringMode> Add<Self> for SmartString<Mode> {

src/ops.rs

@@ -18,6 +18,7 @@ use core::{
 };
 
 pub(crate) trait GenericString: Deref<Target = str> + DerefMut<Target = str> {
+    fn cap(&self) -> usize;
     fn set_size(&mut self, size: usize);
     fn as_mut_capacity_slice(&mut self) -> &mut [u8];
 }
@@ -135,6 +136,15 @@ impl Truncate {
     }
 }
 
+pub(crate) struct EnsureCapacity;
+impl EnsureCapacity {
+    pub(crate) fn cap<S: GenericString>(this: &S, target_cap: usize) -> usize {
+        this.cap().max(target_cap)
+    }
+
+    pub(crate) fn op<S: GenericString>(_this: &mut S, _target_cap: usize) {}
+}
+
 pub(crate) struct Pop;
 impl Pop {
     pub(crate) fn op<S: GenericString>(this: &mut S) -> Option<char> {

src/test.rs

@@ -552,6 +552,42 @@ mod tests {
         assert_eq!(Discriminant::Inline, s.discriminant());
     }
 
+    #[test]
+    fn capacity_modification() {
+        use crate::inline::InlineString;
+
+        const SHORT_TEXT: &'static str = "short enough";
+        static_assertions::const_assert!(SHORT_TEXT.len() < MAX_INLINE - 1);
+        let mut string =
+            SmartString::<Compact>::from_inline(InlineString::try_from(SHORT_TEXT).unwrap());
+
+        string.ensure_capacity(string.len() + 1);
+        assert!(string.capacity() >= SHORT_TEXT.len());
+        assert!(string.is_inline());
+
+        string.ensure_capacity(MAX_INLINE + 1);
+        assert!(!string.is_inline());
+
+        const LARGE_CAPACITY: usize = MAX_INLINE * 40;
+        string.ensure_capacity(LARGE_CAPACITY);
+        assert!(string.capacity() >= LARGE_CAPACITY);
+
+        string.ensure_capacity(LARGE_CAPACITY / 2);
+        assert!(string.capacity() >= LARGE_CAPACITY);
+
+        // Check memory corruptions or size mishandling
+        assert!(string.len() == SHORT_TEXT.len());
+        assert!(string.as_str() == SHORT_TEXT);
+
+        for _ in 0..20 {
+            string = string + "1234567890";
+        }
+        // Read the entire string to check memory corruptions
+        format!("{}", string);
+        assert!(string.capacity() >= LARGE_CAPACITY);
+        assert!(string.len() == SHORT_TEXT.len() + 10 * 20);
+    }
+
     #[test]
     fn from_string() {
         let std_s =