Repository: ParkMyCar/compact_str (Rust)

Issue:
feat: Implement the shrink_to(...) API
This issue for tracking the implementation of the [`shrink_to(...)`](https://doc.rust-lang.org/std/string/struct.String.html#method.shrink_to) method that exists on `String`.
feat: Implement the shrink_to_fit(...) API
This issue for tracking the implementation of the [`shrink_to_fit(...)`](https://doc.rust-lang.org/std/string/struct.String.html#method.shrink_to_fit) method that exists on `String`.


Discussion:

Patch:
diff --git a/compact_str/src/lib.rs b/compact_str/src/lib.rs
index 99e06fc9..405352f8 100644
--- a/compact_str/src/lib.rs
+++ b/compact_str/src/lib.rs
@@ -908,6 +908,73 @@ impl CompactString {
             chars: self[start..end].chars(),
         }
     }
+
+    /// Shrinks the capacity of this [`CompactString`] with a lower bound.
+    ///
+    /// The resulting capactity is never less than the size of 3×[`usize`],
+    /// i.e. the capacity than can be inlined.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// # use compact_str::CompactString;
+    /// let mut s = CompactString::with_capacity(100);
+    /// assert_eq!(s.capacity(), 100);
+    ///
+    /// // if the capacity was already bigger than the argument, the call is a no-op
+    /// s.shrink_to(100);
+    /// assert_eq!(s.capacity(), 100);
+    ///
+    /// s.shrink_to(50);
+    /// assert_eq!(s.capacity(), 50);
+    ///
+    /// // if the string can be inlined, it is
+    /// s.shrink_to(10);
+    /// assert_eq!(s.capacity(), 3 * std::mem::size_of::<usize>());
+    /// ```
+    #[inline]
+    pub fn shrink_to(&mut self, min_capacity: usize) {
+        self.repr.shrink_to(min_capacity);
+    }
+
+    /// Shrinks the capacity of this [`CompactString`] to match its length.
+    ///
+    /// The resulting capactity is never less than the size of 3×[`usize`],
+    /// i.e. the capacity than can be inlined.
+    ///
+    /// This method is effectively the same as calling [`string.shrink_to(0)`].
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// # use compact_str::CompactString;
+    /// let mut s = CompactString::from("This is a string with more than 24 characters.");
+    ///
+    /// s.reserve(100);
+    /// assert!(s.capacity() >= 100);
+    ///
+    ///  s.shrink_to_fit();
+    /// assert_eq!(s.len(), s.capacity());
+    /// ```
+    ///
+    /// ```
+    /// # use compact_str::CompactString;
+    /// let mut s = CompactString::from("short string");
+    ///
+    /// s.reserve(100);
+    /// assert!(s.capacity() >= 100);
+    ///
+    /// s.shrink_to_fit();
+    /// assert_eq!(s.capacity(), 3 * std::mem::size_of::<usize>());
+    /// ```
+    #[inline]
+    pub fn shrink_to_fit(&mut self) {
+        self.repr.shrink_to(0);
+    }
 }
 
 impl Default for CompactString {
diff --git a/compact_str/src/repr/boxed/mod.rs b/compact_str/src/repr/boxed/mod.rs
index 40480a06..803d022c 100644
--- a/compact_str/src/repr/boxed/mod.rs
+++ b/compact_str/src/repr/boxed/mod.rs
@@ -400,7 +400,7 @@ impl BoxString {
     ///
     /// Returns the new capacity as `Ok(…)` if the string was reallocated.
     /// Or the old capacity as `Err(…)` if the string could not be reallocated.
-    fn realloc(&mut self, new_capacity: usize) -> Result<usize, ()> {
+    pub fn realloc(&mut self, new_capacity: usize) -> Result<usize, ()> {
         // can't reallocate a BoxString to a size less than the current length
         if new_capacity < self.len {
             return Err(());
diff --git a/compact_str/src/repr/mod.rs b/compact_str/src/repr/mod.rs
index 4da47dc3..383bdb9a 100644
--- a/compact_str/src/repr/mod.rs
+++ b/compact_str/src/repr/mod.rs
@@ -7,12 +7,11 @@ use std::str::Utf8Error;
 #[cfg(feature = "bytes")]
 mod bytes;
 
-mod iter;
-
 mod boxed;
 mod discriminant;
 mod heap;
 mod inline;
+mod iter;
 mod nonmax;
 mod num;
 mod traits;
@@ -317,6 +316,24 @@ impl Repr {
         // SAFETY: An `ReprUnion` and `Repr` have the same size
         unsafe { std::mem::transmute(self) }
     }
+
+    pub(crate) fn shrink_to(&mut self, min_capacity: usize) {
+        let heap = match self.cast_mut() {
+            MutStrongRepr::Inline(_) => return,
+            MutStrongRepr::Heap(heap) => heap,
+        };
+
+        let old_capacity = heap.string.capacity();
+        let new_capacity = heap.string.len().max(min_capacity);
+        if new_capacity <= MAX_SIZE {
+            // String can be inlined.
+            *self = Repr::from_inline(InlineString::new(heap.string.as_str()));
+        } else if new_capacity < old_capacity {
+            // String can be shrunk.
+            // We can ignore the result. The string keeps its old capacity, but that's okay.
+            let _ = heap.string.realloc(new_capacity);
+        }
+    }
 }
 
 impl Clone for Repr {
diff --git a/fuzz/src/actions.rs b/fuzz/src/actions.rs
index 24b8715f..517caab4 100644
--- a/fuzz/src/actions.rs
+++ b/fuzz/src/actions.rs
@@ -38,6 +38,8 @@ pub enum Action<'a> {
     Drain(u8, u8),
     /// Remove a `char`
     Remove(u8),
+    /// First reserve additional memory, then shrink it
+    ShrinkTo(u16, u16),
 }
 
 impl Action<'_> {
@@ -222,6 +224,26 @@ impl Action<'_> {
                 assert_eq!(control, compact);
                 assert_eq!(control.len(), compact.len());
             }
+            ShrinkTo(a, b) => {
+                let a = (a % 5000) as usize;
+                let b = (b % 5000) as usize;
+                let (reserve, shrink) = (a.max(b), a.min(b));
+
+                control.reserve(reserve);
+                compact.reserve(reserve);
+                assert_eq!(control, compact);
+                assert_eq!(control.len(), compact.len());
+
+                control.shrink_to(shrink);
+                compact.shrink_to(shrink);
+                assert_eq!(control, compact);
+                assert_eq!(control.len(), compact.len());
+
+                control.shrink_to_fit();
+                compact.shrink_to_fit();
+                assert_eq!(control, compact);
+                assert_eq!(control.len(), compact.len());
+            }
         }
     }
 }


Test patch:
diff --git a/compact_str/src/tests.rs b/compact_str/src/tests.rs
index 5a795786..09b59e4d 100644
--- a/compact_str/src/tests.rs
+++ b/compact_str/src/tests.rs
@@ -17,6 +17,8 @@ const MAX_SIZE: usize = 24;
 #[cfg(target_pointer_width = "32")]
 const MAX_SIZE: usize = 12;
 
+const SIXTEEN_MB: usize = 16 * 1024 * 1024;
+
 /// generates random unicode strings, upto 80 chars long
 pub fn rand_unicode() -> impl Strategy<Value = String> {
     proptest::collection::vec(proptest::char::any(), 0..80).prop_map(|v| v.into_iter().collect())
@@ -850,7 +852,6 @@ fn test_into_string_large_string_with_excess_capacity() {
 
 #[test]
 fn test_into_string_where_32_bit_capacity_is_on_heap() {
-    const SIXTEEN_MB: usize = 16 * 1024 * 1024;
     let buf = vec![b'a'; SIXTEEN_MB - 1];
     // SAFETY: `buf` is filled with ASCII `a`s.
     // This primarily speeds up miri, as we don't need to check every byte
@@ -1174,3 +1175,48 @@ fn test_from_utf16() {
         }
     }
 }
+
+#[test]
+fn test_reserve_shrink_roundtrip() {
+    const TEXT: &str = "Hello.";
+
+    let mut s = CompactString::new(TEXT);
+    assert!(!s.is_heap_allocated());
+    assert_eq!(s.capacity(), MAX_SIZE);
+    assert_eq!(s, TEXT);
+
+    s.reserve(128);
+    assert!(s.is_heap_allocated());
+    assert!(s.capacity() >= 128 + TEXT.len());
+    assert_eq!(s, TEXT);
+
+    s.shrink_to(64);
+    assert!(s.is_heap_allocated());
+    assert!(s.capacity() >= 64);
+    assert_eq!(s, TEXT);
+
+    s.shrink_to_fit();
+    assert!(!s.is_heap_allocated());
+    assert_eq!(s.capacity(), MAX_SIZE);
+    assert_eq!(s, TEXT);
+
+    s.reserve(SIXTEEN_MB);
+    assert!(s.is_heap_allocated());
+    assert!(s.capacity() >= SIXTEEN_MB + TEXT.len());
+    assert_eq!(s, TEXT);
+
+    s.shrink_to(64);
+    assert!(s.is_heap_allocated());
+    assert!(s.capacity() >= 64);
+    assert_eq!(s, TEXT);
+
+    s.reserve(SIXTEEN_MB);
+    assert!(s.is_heap_allocated());
+    assert!(s.capacity() >= SIXTEEN_MB + TEXT.len());
+    assert_eq!(s, TEXT);
+
+    s.shrink_to_fit();
+    assert!(!s.is_heap_allocated());
+    assert_eq!(s.capacity(), MAX_SIZE);
+    assert_eq!(s, TEXT);
+}