Repository: ParkMyCar/compact_str
Pull request #37 (2022-01-03 21:25:07)

feature: as_mut_slice API
This PR adds an API to `CompactStr` that provides a mutable slice of the underlying buffer


Patch:
diff --git a/.github/workflows/msrv.yml b/.github/workflows/msrv.yml
index 45f62595..82fe39d7 100644
--- a/.github/workflows/msrv.yml
+++ b/.github/workflows/msrv.yml
@@ -13,18 +13,137 @@ env:
   RUSTFLAGS: "-D warnings"
 
 jobs:
-  test:
-    name: cargo test msrv..
+  test-1_57:
+    name: cargo test 1.57
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
       - uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: stable
-      - name: install cargo hack
-        run: cargo install cargo-hack
-      - name: cargo test msrv..
-        run: |
-          cd compact_str
-          cargo hack test --feature-powerset --optional-deps --version-range 1.49..
+          toolchain: 1.57
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_56:
+    name: cargo test 1.56
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.56
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_55:
+    name: cargo test 1.55
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.55
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_54:
+    name: cargo test 1.54
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.54
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_53:
+    name: cargo test 1.53
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.53
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_52:
+    name: cargo test 1.52
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.52
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_51:
+    name: cargo test 1.51
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.51
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_50:
+    name: cargo test 1.50
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.50.0
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
+
+  test-1_49:
+    name: cargo test 1.49
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: 1.49
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --all-features
diff --git a/compact_str/src/lib.rs b/compact_str/src/lib.rs
index e0f62156..82596313 100644
--- a/compact_str/src/lib.rs
+++ b/compact_str/src/lib.rs
@@ -115,8 +115,11 @@ impl CompactStr {
     /// use compact_str::CompactStr;
     ///
     /// const WORD: usize = std::mem::size_of::<usize>();
-    /// let compact = CompactStr::default();
+    /// let mut compact = CompactStr::default();
     /// assert!(compact.capacity() >= (WORD * 3) - 1);
+    ///
+    /// compact.reserve(200);
+    /// assert!(compact.is_heap_allocated());
     /// ```
     ///
     /// # Panics
@@ -131,6 +134,36 @@ impl CompactStr {
         self.repr.as_str()
     }
 
+    // TODO: Implement a `try_as_mut_slice(...)` that will fail if it results in cloning?
+    //
+    /// Provides a mutable reference to the underlying buffer of bytes.
+    ///
+    /// Note: If the given `CompactStr` is heap allocated, _and_ multiple references exist to the
+    /// underlying buffer (e.g. you previously cloned this `CompactStr`), calling this method will
+    /// clone the entire buffer to prevent silently mutating other owned `CompactStr`s.
+    ///
+    /// ### Futher Explanation
+    /// When a `CompactStr` becomes sufficiently large, the underlying buffer becomes a reference
+    /// counted buffer on the heap. Then, cloning a `CompactStr` increments a reference count
+    /// instead of cloning the entire buffer (very similar to `Arc<str>`). To prevent silently
+    /// mutating the data of other owned `CompactStr`s when taking a mutable slice, we clone the
+    /// underlying buffer and mutate that, if more than one outstanding reference exists.
+    ///
+    /// # Safety
+    /// * All Rust strings, including `CompactStr`, must be valid UTF-8. The caller must guarantee
+    /// that any modifications made to the underlying buffer are valid UTF-8.
+    #[inline]
+    pub unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
+        self.repr.as_mut_slice()
+    }
+
+    /// # Safety
+    /// * TODO: Document safety here
+    #[inline]
+    pub unsafe fn set_len(&mut self, length: usize) {
+        self.repr.set_len(length)
+    }
+
     #[inline]
     pub fn is_heap_allocated(&self) -> bool {
         self.repr.is_heap_allocated()
diff --git a/compact_str/src/repr/heap/arc.rs b/compact_str/src/repr/heap/arc.rs
index b1ae6869..3b1c8ef8 100644
--- a/compact_str/src/repr/heap/arc.rs
+++ b/compact_str/src/repr/heap/arc.rs
@@ -61,6 +61,41 @@ impl ArcString {
         unsafe { str::from_utf8_unchecked(&buffer[..self.len]) }
     }
 
+    #[inline]
+    pub unsafe fn make_mut_slice(&mut self) -> &mut [u8] {
+        if self
+            .inner()
+            .ref_count
+            .compare_exchange(1, 0, Ordering::Acquire, Ordering::Relaxed)
+            .is_err()
+        {
+            // There is more than one reference to this underlying buffer, so we need to make a new
+            // instance and decrement the count of the original by one
+
+            // Make a new instance with the same capacity as self
+            let additional = self.capacity() - self.len();
+            let new = Self::new(self.as_str(), additional);
+
+            // Assign self to our new instsance
+            *self = new;
+        } else {
+            // We were the sole reference of either kind; bump back up the strong ref count.
+            self.inner().ref_count.store(1, Ordering::Release);
+        }
+
+        // Return a mutable slice to the underlying buffer
+        //
+        // SAFETY: If we still have an instance of `ArcString` then we know the pointer to
+        // `ArcStringInner` is valid for at least as long as the provided ref to `self`
+        self.ptr.as_mut().as_mut_bytes()
+    }
+
+    #[inline]
+    pub unsafe fn set_len(&mut self, length: usize) {
+        self.len = length;
+    }
+
+    /// Returns a shared reference to the heap allocated `ArcStringInner`
     #[inline]
     fn inner(&self) -> &ArcStringInner {
         // SAFETY: If we still have an instance of `ArcString` then we know the pointer to
@@ -119,7 +154,7 @@ pub type StrBuffer = [u8; UNKNOWN];
 
 #[repr(C)]
 pub struct ArcStringInner {
-    ref_count: AtomicUsize,
+    pub ref_count: AtomicUsize,
     capacity: usize,
     pub str_buffer: StrBuffer,
 }
@@ -145,6 +180,22 @@ impl ArcStringInner {
         unsafe { slice::from_raw_parts(self.str_buffer.as_ptr(), self.capacity) }
     }
 
+    /// Returns a mutable reference to the underlying buffer of bytes
+    ///
+    /// # Invariants
+    /// * The caller must assert that no other references, or instances of `ArcString` exist before
+    /// calling this method. Otherwise multiple threads could race writing to the underlying buffer.
+    /// * The caller must assert that the underlying buffer is still valid UTF-8
+    #[inline]
+    pub unsafe fn as_mut_bytes(&mut self) -> &mut [u8] {
+        // SAFETY: Since we have an instance of `ArcStringInner` so we know the buffer is still
+        // valid, and we track the capacity with the creation and adjustment of the buffer
+        //
+        // Note: In terms of mutability, it's up to the caller to assert the provided bytes are
+        // value UTF-8
+        slice::from_raw_parts_mut(self.str_buffer.as_mut_ptr(), self.capacity)
+    }
+
     fn layout(capacity: usize) -> alloc::Layout {
         let buffer_layout = alloc::Layout::array::<u8>(capacity).unwrap();
         alloc::Layout::new::<Self>()
diff --git a/compact_str/src/repr/heap/mod.rs b/compact_str/src/repr/heap/mod.rs
index e10bb63d..f0244b9d 100644
--- a/compact_str/src/repr/heap/mod.rs
+++ b/compact_str/src/repr/heap/mod.rs
@@ -38,6 +38,21 @@ impl HeapString {
 
         HeapString { padding, string }
     }
+
+    /// Makes a mutable reference to the underlying buffer, cloning if there is more than one out
+    /// standing reference.
+    ///
+    /// # Invariants
+    /// * Please see `super::Repr` for all invariants
+    #[inline]
+    pub unsafe fn make_mut_slice(&mut self) -> &mut [u8] {
+        self.string.make_mut_slice()
+    }
+
+    #[inline]
+    pub unsafe fn set_len(&mut self, length: usize) {
+        self.string.set_len(length)
+    }
 }
 
 impl From<String> for HeapString {
diff --git a/compact_str/src/repr/inline.rs b/compact_str/src/repr/inline.rs
index 331d7808..5283eddb 100644
--- a/compact_str/src/repr/inline.rs
+++ b/compact_str/src/repr/inline.rs
@@ -87,6 +87,23 @@ impl InlineString {
         // SAFETY: You can only construct an InlineString via a &str
         unsafe { ::std::str::from_utf8_unchecked(slice) }
     }
+
+    /// Provides a mutable reference to the underlying buffer
+    ///
+    /// # Invariants
+    /// * Please see `super::Repr` for all invariants
+    #[inline]
+    pub unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
+        &mut self.buffer[..]
+    }
+
+    #[inline]
+    pub unsafe fn set_len(&mut self, length: usize) {
+        debug_assert!(length <= MAX_INLINE_SIZE);
+
+        // Set the leading bit mask, and then or our length
+        self.metadata = LEADING_BIT_MASK | length as u8;
+    }
 }
 
 static_assertions::assert_eq_size!(InlineString, String);
diff --git a/compact_str/src/repr/mod.rs b/compact_str/src/repr/mod.rs
index ea56a373..b9533b0f 100644
--- a/compact_str/src/repr/mod.rs
+++ b/compact_str/src/repr/mod.rs
@@ -93,7 +93,7 @@ impl Repr {
         let new_capacity = self.len() + additional;
 
         // We already have at least `additional` capacity, so we don't need to do anything
-        if self.capacity() > new_capacity {
+        if self.capacity() >= new_capacity {
             return;
         }
 
@@ -103,7 +103,7 @@ impl Repr {
         // Create a `HeapString` with `text.len() + additional` capacity
         let heap = HeapString::with_additional(self.as_str(), additional);
 
-        // Set self to this new String
+        // Replace `self` with the new Repr
         let heap = ManuallyDrop::new(heap);
         *self = Repr { heap };
     }
@@ -113,6 +113,16 @@ impl Repr {
         self.cast().into_str()
     }
 
+    #[inline]
+    pub unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
+        self.cast_mut().into_mut_slice()
+    }
+
+    #[inline]
+    pub unsafe fn set_len(&mut self, length: usize) {
+        self.cast_mut().set_len(length)
+    }
+
     #[inline]
     pub fn is_heap_allocated(&self) -> bool {
         matches!(self.cast(), StrongRepr::Heap(..))
@@ -141,6 +151,24 @@ impl Repr {
             }
         }
     }
+
+    #[inline(always)]
+    fn cast_mut(&mut self) -> MutStrongRepr<'_> {
+        match self.discriminant() {
+            Discriminant::Heap => {
+                // SAFETY: We checked the discriminant to make sure the union is `heap`
+                MutStrongRepr::Heap(unsafe { &mut self.heap })
+            }
+            Discriminant::Inline => {
+                // SAFETY: We checked the discriminant to make sure the union is `inline`
+                MutStrongRepr::Inline(unsafe { &mut self.inline })
+            }
+            Discriminant::Packed => {
+                // SAFETY: We checked the discriminant to make sure the union is `packed`
+                MutStrongRepr::Packed(unsafe { &mut self.packed })
+            }
+        }
+    }
 }
 
 impl Clone for Repr {
@@ -202,6 +230,33 @@ impl<'a> StrongRepr<'a> {
     }
 }
 
+#[derive(Debug)]
+enum MutStrongRepr<'a> {
+    Heap(&'a mut ManuallyDrop<HeapString>),
+    Inline(&'a mut InlineString),
+    Packed(&'a mut PackedString),
+}
+
+impl<'a> MutStrongRepr<'a> {
+    #[inline]
+    pub unsafe fn into_mut_slice(self) -> &'a mut [u8] {
+        match self {
+            Self::Inline(inline) => inline.as_mut_slice(),
+            Self::Packed(packed) => packed.as_mut_slice(),
+            Self::Heap(heap) => heap.make_mut_slice(),
+        }
+    }
+
+    #[inline]
+    pub unsafe fn set_len(self, length: usize) {
+        match self {
+            Self::Inline(inline) => inline.set_len(length),
+            Self::Packed(packed) => packed.set_len(length),
+            Self::Heap(heap) => heap.set_len(length),
+        }
+    }
+}
+
 assert_eq_size!(Repr, String);
 
 #[cfg(target_pointer_width = "64")]
@@ -265,4 +320,39 @@ mod tests {
         // We should be heap allocated
         assert!(repr.is_heap_allocated());
     }
+
+    #[test]
+    fn test_write_to_buffer() {
+        let mut repr = Repr::new("");
+        let slice = unsafe { repr.as_mut_slice() };
+
+        let word = "abc";
+        let new_len = word.len();
+
+        // write bytes into the `CompactStr`
+        slice[..new_len].copy_from_slice(word.as_bytes());
+        // set the length
+        unsafe { repr.set_len(new_len) }
+
+        assert_eq!(repr.as_str(), word);
+    }
+
+    #[test]
+    fn test_write_to_resized_buffer() {
+        let mut repr = Repr::new("");
+
+        // reserve additional bytes
+        repr.reserve(100);
+        assert!(repr.is_heap_allocated());
+
+        let slice = unsafe { repr.as_mut_slice() };
+
+        let long_word = "hello, I am a very long string that should be allocated on the heap";
+        let new_len = long_word.len();
+
+        slice[..new_len].copy_from_slice(long_word.as_bytes());
+        unsafe { repr.set_len(new_len) }
+
+        assert_eq!(repr.as_str(), long_word);
+    }
 }
diff --git a/compact_str/src/repr/packed.rs b/compact_str/src/repr/packed.rs
index b75369fb..7380c939 100644
--- a/compact_str/src/repr/packed.rs
+++ b/compact_str/src/repr/packed.rs
@@ -62,6 +62,22 @@ impl PackedString {
         // SAFETY: You can only construct a PackedString via a &str
         unsafe { ::std::str::from_utf8_unchecked(&self.buffer) }
     }
+
+    /// Provides a mutable reference to the underlying buffer
+    ///
+    /// # Invariants
+    /// * Please see `super::Repr` for all invariants
+    #[inline]
+    pub unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
+        &mut self.buffer[..]
+    }
+
+    #[inline]
+    pub unsafe fn set_len(&mut self, length: usize) {
+        // An invariant of the Packed representation is that the size always equals MAX_SIZE, so
+        // theres no work to do here, other than assert the length we're trying to set is MAX_SIZE
+        debug_assert_eq!(length, MAX_SIZE);
+    }
 }
 
 static_assertions::assert_eq_size!(PackedString, String);


Test patch:
diff --git a/compact_str/src/tests.rs b/compact_str/src/tests.rs
index d523b12e..19529423 100644
--- a/compact_str/src/tests.rs
+++ b/compact_str/src/tests.rs
@@ -57,6 +57,51 @@ proptest! {
         let word: String = collection.into_iter().collect();
         prop_assert_eq!(&word, &compact);
     }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_reserve_and_write_bytes(word in rand_unicode()) {
+        let mut compact = CompactStr::default();
+        prop_assert!(compact.is_empty());
+
+        // reserve enough space to write our bytes
+        compact.reserve(word.len());
+
+        // SAFETY: We're writing a String which we know is UTF-8
+        let slice = unsafe { compact.as_mut_slice() };
+        slice[..word.len()].copy_from_slice(word.as_bytes());
+
+        // SAFTEY: We know this is the length of our string, since `compact` started with 0 bytes
+        // and we just wrote `word.len()` bytes
+        unsafe { compact.set_len(word.len()) }
+
+        prop_assert_eq!(&word, &compact);
+    }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_reserve_and_write_bytes_allocated_properly(word in rand_unicode()) {
+        let mut compact = CompactStr::default();
+        prop_assert!(compact.is_empty());
+
+        // reserve enough space to write our bytes
+        compact.reserve(word.len());
+
+        // SAFETY: We're writing a String which we know is UTF-8
+        let slice = unsafe { compact.as_mut_slice() };
+        slice[..word.len()].copy_from_slice(word.as_bytes());
+
+        // SAFTEY: We know this is the length of our string, since `compact` started with 0 bytes
+        // and we just wrote `word.len()` bytes
+        unsafe { compact.set_len(word.len()) }
+
+        prop_assert_eq!(compact.len(), word.len());
+
+        // The string should be heap allocated if `word` was >= MAX_INLINED_SIZE
+        //
+        // NOTE: The reserve and write API's don't currently support the Packed representation
+        prop_assert_eq!(compact.is_heap_allocated(), word.len() >= MAX_INLINED_SIZE);
+    }
 }
 
 #[test]