compact_str/repr/
iter.rs

1//! Implementations of the [`FromIterator`] trait to make building [`Repr`]s more ergonomic
2
3use alloc::borrow::Cow;
4use alloc::boxed::Box;
5use alloc::string::String;
6
7use super::{
8    InlineBuffer,
9    Repr,
10    EMPTY,
11    MAX_SIZE,
12};
13use crate::{
14    CompactString,
15    UnwrapWithMsg,
16};
17
18impl FromIterator<char> for Repr {
19    #[inline]
20    fn from_iter<T: IntoIterator<Item = char>>(iter: T) -> Self {
21        let iter = iter.into_iter();
22
23        let (lower_bound, _) = iter.size_hint();
24        let mut this = match Repr::with_capacity(lower_bound) {
25            Ok(this) => this,
26            Err(_) => EMPTY, // Ignore the error and hope that the lower_bound is incorrect.
27        };
28
29        for c in iter {
30            this.push_str(c.encode_utf8(&mut [0; 4]));
31        }
32        this
33    }
34}
35
36impl<'a> FromIterator<&'a char> for Repr {
37    fn from_iter<T: IntoIterator<Item = &'a char>>(iter: T) -> Self {
38        iter.into_iter().copied().collect()
39    }
40}
41
42fn from_as_ref_str_iterator<S, I>(mut iter: I) -> Repr
43where
44    S: AsRef<str>,
45    I: Iterator<Item = S>,
46    String: core::iter::Extend<S>,
47    String: FromIterator<S>,
48{
49    // Note: We don't check the lower bound here like we do in the character iterator because it's
50    // possible for the iterator to be full of empty strings! In which case checking the lower bound
51    // could cause us to heap allocate when there's no need.
52
53    // Continuously pull strings from the iterator
54    let mut curr_len = 0;
55    let mut inline_buf = InlineBuffer::new_const("");
56    while let Some(s) = iter.next() {
57        let str_slice = s.as_ref();
58        let bytes_len = str_slice.len();
59
60        // this new string is too large to fit into our inline buffer, so heap allocate the rest
61        if bytes_len + curr_len > MAX_SIZE {
62            let (min_remaining, _) = iter.size_hint();
63            let mut string = String::with_capacity(bytes_len + curr_len + min_remaining);
64
65            // push existing strings onto the heap
66            // SAFETY: `inline_buf` has been filled with `&str`s which are valid UTF-8
67            string.push_str(unsafe { core::str::from_utf8_unchecked(&inline_buf.0[..curr_len]) });
68            // push current string onto the heap
69            string.push_str(str_slice);
70            // extend heap with remaining strings
71            string.extend(iter);
72
73            return Repr::from_string(string, true).unwrap_with_msg();
74        }
75
76        // write the current string into a slice of the unoccupied space
77        inline_buf.0[curr_len..][..bytes_len].copy_from_slice(str_slice.as_bytes());
78        curr_len += bytes_len;
79    }
80
81    // SAFETY: Everything we just pushed onto the buffer is a `str` which is valid UTF-8
82    unsafe { inline_buf.set_len(curr_len) }
83
84    Repr::from_inline(inline_buf)
85}
86
87impl<'a> FromIterator<&'a str> for Repr {
88    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
89        from_as_ref_str_iterator(iter.into_iter())
90    }
91}
92
93impl FromIterator<Box<str>> for Repr {
94    fn from_iter<T: IntoIterator<Item = Box<str>>>(iter: T) -> Self {
95        from_as_ref_str_iterator(iter.into_iter())
96    }
97}
98
99impl FromIterator<String> for Repr {
100    fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
101        from_as_ref_str_iterator(iter.into_iter())
102    }
103}
104
105impl FromIterator<CompactString> for Repr {
106    fn from_iter<T: IntoIterator<Item = CompactString>>(iter: T) -> Self {
107        from_as_ref_str_iterator(iter.into_iter())
108    }
109}
110
111impl<'a> FromIterator<Cow<'a, str>> for Repr {
112    fn from_iter<T: IntoIterator<Item = Cow<'a, str>>>(iter: T) -> Self {
113        from_as_ref_str_iterator(iter.into_iter())
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use alloc::string::String;
120
121    use super::Repr;
122
123    #[test]
124    fn short_char_iter() {
125        let chars = ['a', 'b', 'c'];
126        let repr: Repr = chars.iter().collect();
127
128        assert_eq!(repr.as_str(), "abc");
129        assert!(!repr.is_heap_allocated());
130    }
131
132    #[test]
133    fn short_char_ref_iter() {
134        let chars = ['a', 'b', 'c'];
135        let repr: Repr = chars.iter().collect();
136
137        assert_eq!(repr.as_str(), "abc");
138        assert!(!repr.is_heap_allocated());
139    }
140
141    #[test]
142    #[cfg_attr(target_pointer_width = "32", ignore)]
143    fn packed_char_iter() {
144        let chars = [
145            '\u{92f01}',
146            '\u{81515}',
147            '\u{81515}',
148            '\u{81515}',
149            '\u{81515}',
150            '\u{41515}',
151        ];
152
153        let repr: Repr = chars.iter().collect();
154        let s: String = chars.iter().collect();
155
156        assert_eq!(repr.as_str(), s.as_str());
157        assert!(!repr.is_heap_allocated());
158    }
159
160    #[test]
161    fn long_char_iter() {
162        let long = "This is supposed to be a really long string";
163        let repr: Repr = long.chars().collect();
164
165        assert_eq!(repr.as_str(), "This is supposed to be a really long string");
166        assert!(repr.is_heap_allocated());
167    }
168
169    #[test]
170    fn short_string_iter() {
171        let strings = vec!["hello", "world"];
172        let repr: Repr = strings.into_iter().collect();
173
174        assert_eq!(repr.as_str(), "helloworld");
175        assert!(!repr.is_heap_allocated());
176    }
177
178    #[test]
179    fn long_short_string_iter() {
180        let strings = vec![
181            "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16",
182            "17", "18", "19", "20",
183        ];
184        let repr: Repr = strings.into_iter().collect();
185
186        assert_eq!(repr.as_str(), "1234567891011121314151617181920");
187        assert!(repr.is_heap_allocated());
188    }
189}