Repository: ParkMyCar/compact_str
Pull request #194 (2022-08-06 22:54:19)

feat: from_utf8_unchecked API
This PR adds `CompactString::from_utf8_unchecked(...)`, which has the same behaviors as [`String::from_utf8_unchecked(...)`](https://doc.rust-lang.org/std/string/struct.String.html#method.from_utf8_unchecked). We also add unit tests (which get Miri coverage), proptests, and update our fuzzing harness.

Fixes: #138 


Patch:
diff --git a/compact_str/src/lib.rs b/compact_str/src/lib.rs
index 5ebca697..c25d365e 100644
--- a/compact_str/src/lib.rs
+++ b/compact_str/src/lib.rs
@@ -253,6 +253,40 @@ impl CompactString {
         Ok(CompactString { repr })
     }
 
+    /// Converts a vector of bytes to a [`CompactString`] without checking that the string contains
+    /// valid UTF-8.
+    ///
+    /// See the safe version, [`CompactString::from_utf8`], for more details.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe because it does not check that the bytes passed to it are valid
+    /// UTF-8. If this constraint is violated, it may cause memory unsafety issues with future users
+    /// of the [`CompactString`], as the rest of the standard library assumes that
+    /// [`CompactString`]s are valid UTF-8.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// # use compact_str::CompactString;
+    /// // some bytes, in a vector
+    /// let sparkle_heart = vec![240, 159, 146, 150];
+    ///
+    /// let sparkle_heart = unsafe {
+    ///     CompactString::from_utf8_unchecked(sparkle_heart)
+    /// };
+    ///
+    /// assert_eq!("💖", sparkle_heart);
+    /// ```
+    #[inline]
+    #[must_use]
+    pub unsafe fn from_utf8_unchecked<B: AsRef<[u8]>>(buf: B) -> Self {
+        let repr = Repr::from_utf8_unchecked(buf);
+        CompactString { repr }
+    }
+
     /// Decode a [`UTF-16`](https://en.wikipedia.org/wiki/UTF-16) slice of bytes into a
     /// [`CompactString`], returning an [`Err`] if the slice contains any invalid data.
     ///
diff --git a/compact_str/src/repr/mod.rs b/compact_str/src/repr/mod.rs
index 383bdb9a..f0716529 100644
--- a/compact_str/src/repr/mod.rs
+++ b/compact_str/src/repr/mod.rs
@@ -111,6 +111,40 @@ impl Repr {
         Ok(Self::new(s))
     }
 
+    #[inline]
+    pub unsafe fn from_utf8_unchecked<B: AsRef<[u8]>>(buf: B) -> Self {
+        let bytes = buf.as_ref();
+        let bytes_len = bytes.len();
+
+        // Create a Repr with enough capacity for the entire buffer
+        let mut repr = Repr::with_capacity(bytes_len);
+
+        // There's an edge case where the final byte of this buffer == `HEAP_MASK`, which is
+        // invalid UTF-8, but would result in us creating an inline variant, that identifies as
+        // a heap variant. If a user ever tried to reference the data at all, we'd incorrectly
+        // try and read data from an invalid memory address, causing undefined behavior.
+        if bytes_len == MAX_SIZE {
+            let last_byte = bytes[bytes_len - 1];
+            // If we hit the edge case, reserve additional space to make the repr becomes heap
+            // allocated, which prevents us from writing this last byte inline
+            if last_byte >= 0b11000000 {
+                repr.reserve(MAX_SIZE + 1);
+            }
+        }
+
+        // SAFETY: The caller is responsible for making sure the provided buffer is UTF-8. This
+        // invariant is documented in the public API
+        let slice = repr.as_mut_slice();
+        // write the chunk into the Repr
+        slice[..bytes_len].copy_from_slice(bytes);
+
+        // Set the length of the Repr
+        // SAFETY: We just wrote the entire `buf` into the Repr
+        repr.set_len(bytes_len);
+
+        repr
+    }
+
     #[inline]
     pub fn from_string(s: String) -> Self {
         if s.capacity() == 0 {
@@ -557,6 +591,7 @@ crate::asserts::assert_size!(Repr, 12);
 mod tests {
     use super::{
         Repr,
+        HEAP_MASK,
         MAX_SIZE,
     };
 
@@ -810,4 +845,15 @@ mod tests {
         assert!(!repr.is_heap_allocated());
         assert_eq!(repr.as_str(), "hello world!");
     }
+
+    #[test]
+    fn test_from_utf8_unchecked_heap_mask() {
+        let bytes = [HEAP_MASK; MAX_SIZE];
+
+        // this isn't a valid string, but we should be able to read back the underlying buffer
+        let repr = unsafe { Repr::from_utf8_unchecked(bytes) };
+
+        assert_eq!(repr.len(), MAX_SIZE);
+        assert_eq!(repr.as_slice(), bytes);
+    }
 }
diff --git a/fuzz/src/creation.rs b/fuzz/src/creation.rs
index a41711d6..c16364bf 100644
--- a/fuzz/src/creation.rs
+++ b/fuzz/src/creation.rs
@@ -22,6 +22,8 @@ use crate::MAX_INLINE_LENGTH;
 pub enum Creation<'a> {
     /// Create using [`CompactString::from_utf8`]
     Bytes(&'a [u8]),
+    /// Create using [`CompactString::from_utf8_unchecked`]
+    BytesUnchecked(&'a [u8]),
     /// Create using [`CompactString::from_utf16`]
     BytesUtf16(Vec<u16>),
     /// Create using [`CompactString::from_utf8_buf`]
@@ -285,6 +287,44 @@ impl Creation<'_> {
                     _ => panic!("CompactString and core::str read UTF-8 differently?"),
                 }
             }
+            BytesUnchecked(data) => {
+                // The data provided might not be valid UTF-8. We mainly want to make sure we don't
+                // panic, and the data is written correctly. Before returning either of these types
+                // we'll make sure they contain valid data
+                let compact = unsafe { CompactString::from_utf8_unchecked(data) };
+                let std_str = unsafe { String::from_utf8_unchecked(data.to_vec()) };
+
+                // make sure our data didn't somehow get longer
+                assert_eq!(data.len(), compact.len());
+                assert_eq!(compact.len(), std_str.len());
+
+                // make sure the data written is the same
+                assert_eq!(compact.as_bytes(), std_str.as_bytes());
+
+                let data_is_valid = std::str::from_utf8(data);
+                let compact_is_valid = std::str::from_utf8(compact.as_bytes());
+                let std_str_is_valid = std::str::from_utf8(std_str.as_bytes());
+
+                match (data_is_valid, compact_is_valid, std_str_is_valid) {
+                    (Ok(d), Ok(c), Ok(s)) => {
+                        // if we get &str's back, make sure they're all equal
+                        assert_eq!(d, c);
+                        assert_eq!(c, s);
+
+                        // we have valid UTF-8 data! we can return a pair
+                        Some((compact, std_str))
+                    }
+                    (Err(d), Err(c), Err(s)) => {
+                        // if we get errors back, the errors should be the same
+                        assert_eq!(d, c);
+                        assert_eq!(c, s);
+
+                        // we don't have valid UTF-8 data, so we can't return anything
+                        None
+                    }
+                    _ => panic!("data, CompactString, and String disagreed?"),
+                }
+            }
             BytesUtf16(data) => {
                 let compact = CompactString::from_utf16(&data);
                 let std_str = String::from_utf16(&data);


Test patch:
diff --git a/compact_str/src/tests.rs b/compact_str/src/tests.rs
index 09b59e4d..0b1be1b2 100644
--- a/compact_str/src/tests.rs
+++ b/compact_str/src/tests.rs
@@ -287,6 +287,40 @@ fn proptest_remove(#[strategy(rand_unicode_with_range(1..80))] mut control: Stri
     }
 }
 
+#[proptest]
+#[cfg_attr(miri, ignore)]
+fn proptest_from_utf8_unchecked(#[strategy(rand_bytes())] bytes: Vec<u8>) {
+    let compact = unsafe { CompactString::from_utf8_unchecked(&bytes) };
+    let std_str = unsafe { String::from_utf8_unchecked(bytes.clone()) };
+
+    // we might not make valid strings, but we should be able to read the underlying bytes
+    assert_eq!(compact.as_bytes(), std_str.as_bytes());
+    assert_eq!(compact.as_bytes(), bytes);
+
+    // make sure the length is correct
+    assert_eq!(compact.len(), bytes.len());
+
+    // check if we were valid UTF-8, if so, assert the data written into the CompactString is
+    // correct
+    let data_is_valid = std::str::from_utf8(&bytes);
+    let compact_is_valid = std::str::from_utf8(compact.as_bytes());
+    let std_str_is_valid = std::str::from_utf8(std_str.as_bytes());
+
+    match (data_is_valid, compact_is_valid, std_str_is_valid) {
+        (Ok(d), Ok(c), Ok(s)) => {
+            // if we get &str's back, make sure they're all equal
+            assert_eq!(d, c);
+            assert_eq!(c, s);
+        }
+        (Err(d), Err(c), Err(s)) => {
+            // if we get errors back, the errors should be the same
+            assert_eq!(d, c);
+            assert_eq!(c, s);
+        }
+        _ => panic!("data, CompactString, and String disagreed?"),
+    }
+}
+
 #[test]
 fn test_const_creation() {
     const EMPTY: CompactString = CompactString::new_inline("");
@@ -1220,3 +1254,38 @@ fn test_reserve_shrink_roundtrip() {
     assert_eq!(s.capacity(), MAX_SIZE);
     assert_eq!(s, TEXT);
 }
+
+#[test]
+fn test_from_utf8_unchecked_sanity() {
+    let text = "hello 🌎, you are nice";
+    let compact = unsafe { CompactString::from_utf8_unchecked(text) };
+
+    assert_eq!(compact, text);
+}
+
+#[test]
+fn test_from_utf8_unchecked_long() {
+    let bytes = [255; 2048];
+    let compact = unsafe { CompactString::from_utf8_unchecked(bytes) };
+
+    assert_eq!(compact.len(), 2048);
+    assert_eq!(compact.as_bytes(), bytes);
+}
+
+#[test]
+fn test_from_utf8_unchecked_short() {
+    let bytes = [255; 10];
+    let compact = unsafe { CompactString::from_utf8_unchecked(bytes) };
+
+    assert_eq!(compact.len(), 10);
+    assert_eq!(compact.as_bytes(), bytes);
+}
+
+#[test]
+fn test_from_utf8_unchecked_empty() {
+    let bytes = [255; 0];
+    let compact = unsafe { CompactString::from_utf8_unchecked(bytes) };
+
+    assert_eq!(compact.len(), 0);
+    assert_eq!(compact.as_bytes(), bytes);
+}