unarray/
from_iter.rs

1use crate::{mark_initialized, uninit_buf};
2use core::iter::FromIterator;
3
4/// A wrapper type to collect an [`Iterator`] into an array
5///
6/// ```
7/// # use unarray::*;
8/// let iter = vec![1, 2, 3].into_iter();
9/// let ArrayFromIter(array) = iter.collect();
10///
11/// assert_eq!(array, Some([1, 2, 3]));
12/// ```
13/// Since iterators don't carry compile-time information about their length (even
14/// [`core::iter::ExactSizeIterator`] only provides this at runtime), collecting may fail if the
15/// iterator doesn't yield **exactly** `N` elements:
16/// ```
17/// use unarray::*;
18/// let too_many = vec![1, 2, 3, 4].into_iter();
19/// let ArrayFromIter::<i32, 3>(option) = too_many.collect();
20/// assert!(option.is_none());
21///
22/// let too_few = vec![1, 2].into_iter();
23/// let ArrayFromIter::<i32, 3>(option) = too_few.collect();
24/// assert!(option.is_none());
25/// ```
26pub struct ArrayFromIter<T, const N: usize>(pub Option<[T; N]>);
27
28impl<T, const N: usize> FromIterator<T> for ArrayFromIter<T, N> {
29    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
30        let mut buffer = uninit_buf::<T, N>();
31        let mut iter = iter.into_iter();
32        let mut buf_iter = buffer.iter_mut();
33        let mut num_written = 0;
34
35        loop {
36            let item = iter.next();
37            let slot = buf_iter.next();
38
39            match (item, slot) {
40                (Some(item), Some(slot)) => {
41                    num_written += 1;
42                    slot.write(item);
43                }
44                // error case, we should free the previous ones and continue
45                (Some(_), None) | (None, Some(_)) => {
46                    // SAFETY:
47                    // There are `num_written` elements fully initialized, so we can safely drop
48                    // them
49                    buffer
50                        .iter_mut()
51                        .take(num_written)
52                        .for_each(|slot| unsafe { slot.assume_init_drop() });
53
54                    return Self(None);
55                }
56                // SAFETY
57                // If this is reached, every prior iteration of the loop has matched
58                // (Some(_), Some(_)). As such, both iterators have yielded the same number of
59                // elements, so every slot has been written to
60                (None, None) => return Self(Some(unsafe { mark_initialized(buffer) })),
61            };
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68
69    use core::{
70        convert::TryInto,
71        sync::atomic::{AtomicUsize, Ordering},
72    };
73
74    use crate::testing::vec_strategy;
75    use proptest::{prop_assert, prop_assert_eq};
76    use test_strategy::proptest;
77
78    use super::*;
79
80    #[test]
81    fn can_collect_array_from_iter() {
82        let iter = vec![1, 2, 3].into_iter();
83
84        let ArrayFromIter(array) = iter.collect();
85        assert_eq!(array.unwrap(), [1, 2, 3]);
86    }
87
88    #[test]
89    fn fails_if_incorrect_number_of_elements() {
90        let iter = [1, 2, 3].iter();
91        let ArrayFromIter::<_, 4>(array) = iter.collect();
92        assert!(array.is_none());
93
94        let iter = [1, 2, 3].iter();
95        let ArrayFromIter::<_, 2>(array) = iter.collect();
96        assert!(array.is_none());
97    }
98
99    const LEN: usize = 100;
100    const SHORT_LEN: usize = LEN - 1;
101    const LONG_LEN: usize = LEN + 1;
102
103    #[derive(Clone)]
104    struct IncrementOnDrop<'a>(&'a AtomicUsize);
105    impl Drop for IncrementOnDrop<'_> {
106        fn drop(&mut self) {
107            self.0.fetch_add(1, Ordering::Relaxed);
108        }
109    }
110
111    #[test]
112    fn doesnt_leak_too_long() {
113        let drop_count = 0.into();
114        let ArrayFromIter::<_, 3>(_) = vec![IncrementOnDrop(&drop_count); 4].into_iter().collect();
115        // since it failed, all 4 should be dropped
116        assert_eq!(drop_count.load(Ordering::Relaxed), 4);
117    }
118
119    #[test]
120    fn doesnt_leak_too_short() {
121        let drop_count = 0.into();
122        let ArrayFromIter::<_, 3>(_) = vec![IncrementOnDrop(&drop_count); 2].into_iter().collect();
123        // since it failed, both should be dropped
124        assert_eq!(drop_count.load(Ordering::Relaxed), 2);
125    }
126
127    #[proptest]
128    #[cfg_attr(miri, ignore)]
129    fn undersized_proptest(#[strategy(vec_strategy(LEN))] vec: Vec<String>) {
130        let ArrayFromIter::<String, SHORT_LEN>(array) = vec.into_iter().collect();
131        prop_assert!(array.is_none());
132    }
133
134    #[proptest]
135    #[cfg_attr(miri, ignore)]
136    fn oversized_proptest(#[strategy(vec_strategy(LEN))] vec: Vec<String>) {
137        let ArrayFromIter::<String, LONG_LEN>(array) = vec.into_iter().collect();
138        prop_assert!(array.is_none());
139    }
140
141    #[proptest]
142    #[cfg_attr(miri, ignore)]
143    fn just_right_proptest(#[strategy(vec_strategy(LEN))] vec: Vec<String>) {
144        let expected: [String; LEN] = vec.clone().try_into().unwrap();
145        let ArrayFromIter(array) = vec.into_iter().collect();
146        prop_assert_eq!(array.unwrap(), expected);
147    }
148}
unarray/from_iter.rs

unarray/
from_iter.rs
