Repository: ParkMyCar/compact_str
Pull request #23 (2021-11-14 22:20:21)

impl FromIterator<...> for CompactStr
This PR implements the `FromIterator` trait for a common string types to make building `CompactStr`s more ergonomic


Patch:
diff --git a/compact_str/src/lib.rs b/compact_str/src/lib.rs
index 002e399f..02cc5740 100644
--- a/compact_str/src/lib.rs
+++ b/compact_str/src/lib.rs
@@ -16,6 +16,7 @@ use core::{
     cmp::Ordering,
     fmt,
     hash::{Hash, Hasher},
+    iter::FromIterator,
     ops::Deref,
     str::FromStr,
 };
@@ -201,4 +202,39 @@ impl fmt::Display for CompactStr {
     }
 }
 
+impl FromIterator<char> for CompactStr {
+    fn from_iter<T: IntoIterator<Item = char>>(iter: T) -> Self {
+        let repr = iter.into_iter().collect();
+        CompactStr { repr }
+    }
+}
+
+impl<'a> FromIterator<&'a char> for CompactStr {
+    fn from_iter<T: IntoIterator<Item = &'a char>>(iter: T) -> Self {
+        let repr = iter.into_iter().collect();
+        CompactStr { repr }
+    }
+}
+
+impl<'a> FromIterator<&'a str> for CompactStr {
+    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
+        let repr = iter.into_iter().collect();
+        CompactStr { repr }
+    }
+}
+
+impl FromIterator<Box<str>> for CompactStr {
+    fn from_iter<T: IntoIterator<Item = Box<str>>>(iter: T) -> Self {
+        let repr = iter.into_iter().collect();
+        CompactStr { repr }
+    }
+}
+
+impl FromIterator<String> for CompactStr {
+    fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
+        let repr = iter.into_iter().collect();
+        CompactStr { repr }
+    }
+}
+
 static_assertions::assert_eq_size!(CompactStr, String);
diff --git a/compact_str/src/repr/heap.rs b/compact_str/src/repr/heap.rs
index 60e979a6..a15e988e 100644
--- a/compact_str/src/repr/heap.rs
+++ b/compact_str/src/repr/heap.rs
@@ -21,4 +21,13 @@ impl HeapString {
     }
 }
 
+impl From<String> for HeapString {
+    fn from(s: String) -> Self {
+        let padding = PADDING;
+        let string = s.into();
+
+        HeapString { padding, string }
+    }
+}
+
 static_assertions::assert_eq_size!(HeapString, String);
diff --git a/compact_str/src/repr/inline.rs b/compact_str/src/repr/inline.rs
index 500290c2..40f49c59 100644
--- a/compact_str/src/repr/inline.rs
+++ b/compact_str/src/repr/inline.rs
@@ -58,6 +58,13 @@ impl InlineString {
         InlineString { metadata, buffer }
     }
 
+    /// Creates an `InlineString` from raw parts without checking that it's valid UTF-8.
+    #[inline]
+    pub const unsafe fn from_parts(len: usize, buffer: [u8; MAX_INLINE_SIZE]) -> Self {
+        let metadata = (len as u8) | LEADING_BIT_MASK;
+        InlineString { metadata, buffer }
+    }
+
     #[inline]
     pub const fn len(&self) -> usize {
         (self.metadata & !LEADING_BIT_MASK) as usize
diff --git a/compact_str/src/repr/iter.rs b/compact_str/src/repr/iter.rs
new file mode 100644
index 00000000..265ee836
--- /dev/null
+++ b/compact_str/src/repr/iter.rs
@@ -0,0 +1,187 @@
+//! Implementations of the `FromIterator` trait to make building `CompactStr`s more ergonomic
+
+use super::{inline::MAX_INLINE_SIZE, HeapString, InlineString, Repr};
+use core::{iter::FromIterator, mem::ManuallyDrop};
+
+impl FromIterator<char> for Repr {
+    fn from_iter<T: IntoIterator<Item = char>>(iter: T) -> Self {
+        let mut iter = iter.into_iter();
+
+        // If the size hint indicates we can't store this inline, then create a heap string
+        let (size_hint, _) = iter.size_hint();
+        if size_hint > MAX_INLINE_SIZE {
+            let heap = HeapString::from(iter.collect::<String>());
+            return Repr {
+                heap: ManuallyDrop::new(heap),
+            };
+        }
+
+        // Otherwise, continuously pull chars from the iterator
+        let mut curr_len = 0;
+        let mut inline_buf = [0u8; MAX_INLINE_SIZE];
+        while let Some(c) = iter.next() {
+            let char_len = c.len_utf8();
+
+            // If this new character is too large to fit into the inline buffer, then create a heap string
+            if char_len + curr_len > MAX_INLINE_SIZE {
+                let (min_remaining, _) = iter.size_hint();
+                let mut heap_buf = String::with_capacity(char_len + curr_len + min_remaining);
+
+                // push existing characters onto the heap
+                // SAFETY: `inline_buf` has been filled with `char`s which are valid UTF-8
+                heap_buf
+                    .push_str(unsafe { core::str::from_utf8_unchecked(&inline_buf[..curr_len]) });
+                // push current char onto the heap
+                heap_buf.push(c);
+                // extend heap with remaining characters
+                heap_buf.extend(iter);
+
+                let heap = HeapString::from(heap_buf);
+                return Repr {
+                    heap: ManuallyDrop::new(heap),
+                };
+            }
+
+            // write the current char into a slice of the unoccupied space
+            c.encode_utf8(&mut inline_buf[curr_len..]);
+            curr_len += char_len;
+        }
+
+        // TODO: Support PackedString here in an efficient way
+
+        // SAFETY: We know `inline_buf` is valid UTF-8 because it consists entriely of `char`s
+        let inline = unsafe { InlineString::from_parts(curr_len, inline_buf) };
+        Repr { inline }
+    }
+}
+
+impl<'a> FromIterator<&'a char> for Repr {
+    fn from_iter<T: IntoIterator<Item = &'a char>>(iter: T) -> Self {
+        iter.into_iter().copied().collect()
+    }
+}
+
+fn from_as_ref_str_iterator<S, I>(mut iter: I) -> Repr
+where
+    S: AsRef<str>,
+    I: Iterator<Item = S>,
+    String: core::iter::Extend<S>,
+    String: FromIterator<S>,
+{
+    // If there are more strings than we can fit bytes inline, then immediately make a `HeapString`
+    let (size_hint, _) = iter.size_hint();
+    if size_hint > MAX_INLINE_SIZE {
+        let heap = HeapString::from(iter.collect::<String>());
+        return Repr {
+            heap: ManuallyDrop::new(heap),
+        };
+    }
+
+    // Otherwise, continuously pull strings from the iterator
+    let mut curr_len = 0;
+    let mut inline_buf = [0u8; MAX_INLINE_SIZE];
+    while let Some(s) = iter.next() {
+        let str_slice = s.as_ref();
+        let bytes_len = str_slice.len();
+
+        // this new string is too large to fit into our inline buffer, so heap allocate the rest
+        if bytes_len + curr_len > MAX_INLINE_SIZE {
+            let (min_remaining, _) = iter.size_hint();
+            let mut heap_buf = String::with_capacity(bytes_len + curr_len + min_remaining);
+
+            // push existing strings onto the heap
+            // SAFETY: `inline_buf` has been filled with `&str`s which are valid UTF-8
+            heap_buf.push_str(unsafe { core::str::from_utf8_unchecked(&inline_buf[..curr_len]) });
+            // push current string onto the heap
+            heap_buf.push_str(str_slice);
+            // extend heap with remaining strings
+            heap_buf.extend(iter);
+
+            let heap = HeapString::from(heap_buf);
+            return Repr {
+                heap: ManuallyDrop::new(heap),
+            };
+        }
+
+        // write the current string into a slice of the unoccupied space
+        (&mut inline_buf[curr_len..][..bytes_len]).copy_from_slice(str_slice.as_bytes());
+        curr_len += bytes_len;
+    }
+
+    // TODO: Support PackedString here in an efficient way
+
+    // SAFETY: We know `inline_buf` is valid UTF-8 because it consists entriely of `&str`s
+    let inline = unsafe { InlineString::from_parts(curr_len, inline_buf) };
+    Repr { inline }
+}
+
+impl<'a> FromIterator<&'a str> for Repr {
+    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
+        from_as_ref_str_iterator(iter.into_iter())
+    }
+}
+
+impl FromIterator<Box<str>> for Repr {
+    fn from_iter<T: IntoIterator<Item = Box<str>>>(iter: T) -> Self {
+        from_as_ref_str_iterator(iter.into_iter())
+    }
+}
+
+impl FromIterator<String> for Repr {
+    fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
+        from_as_ref_str_iterator(iter.into_iter())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Repr;
+
+    #[test]
+    fn short_char_iter() {
+        let chars = ['a', 'b', 'c'];
+        let repr: Repr = chars.into_iter().collect();
+
+        assert_eq!(repr.as_str(), "abc");
+        assert!(!repr.is_heap_allocated());
+    }
+
+    #[test]
+    fn short_char_ref_iter() {
+        let chars = ['a', 'b', 'c'];
+        let repr: Repr = chars.iter().collect();
+
+        assert_eq!(repr.as_str(), "abc");
+        assert!(!repr.is_heap_allocated());
+    }
+
+    #[test]
+    fn long_char_iter() {
+        let long = "This is supposed to be a really long string";
+        let repr: Repr = long.chars().collect();
+
+        assert_eq!(repr.as_str(), "This is supposed to be a really long string");
+        assert!(repr.is_heap_allocated());
+    }
+
+    #[test]
+    fn short_string_iter() {
+        let strings = ["hello", "world"];
+        let repr: Repr = strings.into_iter().collect();
+
+        assert_eq!(repr.as_str(), "helloworld");
+        assert!(!repr.is_heap_allocated());
+    }
+
+    #[test]
+    fn long_short_string_iter() {
+        let strings = [
+            "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16",
+            "17", "18", "19", "20",
+        ];
+        let repr: Repr = strings.into_iter().collect();
+
+        assert_eq!(repr.as_str(), "1234567891011121314151617181920");
+        assert!(repr.is_heap_allocated());
+    }
+}
diff --git a/compact_str/src/repr/mod.rs b/compact_str/src/repr/mod.rs
index 224e8643..6ae88fbe 100644
--- a/compact_str/src/repr/mod.rs
+++ b/compact_str/src/repr/mod.rs
@@ -1,6 +1,8 @@
 use static_assertions::{assert_eq_size, const_assert_eq};
 use std::mem::ManuallyDrop;
 
+mod iter;
+
 mod discriminant;
 mod heap;
 mod inline;


Test patch:
diff --git a/compact_str/src/tests.rs b/compact_str/src/tests.rs
index a4181529..9b8fb0f6 100644
--- a/compact_str/src/tests.rs
+++ b/compact_str/src/tests.rs
@@ -12,6 +12,11 @@ fn rand_unicode() -> impl Strategy<Value = String> {
     proptest::collection::vec(proptest::char::any(), 0..80).prop_map(|v| v.into_iter().collect())
 }
 
+// generates groups upto 40 strings long of random unicode strings, upto 80 chars long
+fn rand_unicode_collection() -> impl Strategy<Value = Vec<String>> {
+    proptest::collection::vec(rand_unicode(), 0..40)
+}
+
 proptest! {
     #[test]
     fn test_strings_roundtrip(word in rand_unicode()) {
@@ -32,6 +37,19 @@ proptest! {
             prop_assert!(compact.is_heap_allocated())
         }
     }
+
+    #[test]
+    fn test_char_iterator_roundtrips(word in rand_unicode()) {
+        let compact: CompactStr = word.clone().chars().collect();
+        prop_assert_eq!(&word, &compact)
+    }
+
+    #[test]
+    fn test_string_iterator_roundtrips(collection in rand_unicode_collection()) {
+        let compact: CompactStr = collection.clone().into_iter().collect();
+        let word: String = collection.into_iter().collect();
+        prop_assert_eq!(&word, &compact);
+    }
 }
 
 #[test]
@@ -140,3 +158,13 @@ fn test_from_str_trait() {
 
     assert_eq!(s, c);
 }
+
+#[test]
+fn test_from_char_iter() {
+    let s = "\u{0} 0 \u{0}a𐀀𐀀 𐀀a𐀀";
+    println!("{}", s.len());
+    let compact: CompactStr = s.chars().into_iter().collect();
+
+    assert!(compact.is_heap_allocated());
+    assert_eq!(s, compact);
+}