ring/polyfill/
array_flat_map.rs

Help
1use core::iter::FlatMap;
2
3/// A specialized version of `core::iter::FlatMap` for mapping over exact-sized
4/// iterators with a function that returns an array.
5///
6/// `ArrayFlatMap` differs from `FlatMap` in that `ArrayFlatMap` implements
7/// `ExactSizeIterator`. Since the result of `F` always has `LEN` elements, if
8/// `I` is an exact-sized iterator of length `inner_len` then we know the
9/// length of the flat-mapped result is `inner_len * LEN`. (The constructor
10/// verifies that this multiplication doesn't overflow `usize`.)
11#[derive(Clone)]
12pub struct ArrayFlatMap<I, Item, F, const LEN: usize> {
13    inner: FlatMap<I, [Item; LEN], F>,
14    remaining: usize,
15}
16
17impl<I, Item, F, const LEN: usize> ArrayFlatMap<I, Item, F, LEN>
18where
19    I: ExactSizeIterator,
20    F: FnMut(I::Item) -> [Item; LEN],
21{
22    /// Constructs an `ArrayFlatMap` wrapping the given iterator, using the
23    /// given function
24    pub fn new(inner: I, f: F) -> Option<Self> {
25        let remaining = inner.len().checked_mul(LEN)?;
26        let inner = inner.flat_map(f);
27        Some(Self { inner, remaining })
28    }
29}
30
31impl<I, Item, F, const LEN: usize> Iterator for ArrayFlatMap<I, Item, F, LEN>
32where
33    I: Iterator,
34    F: FnMut(I::Item) -> [Item; LEN],
35{
36    type Item = Item;
37
38    fn next(&mut self) -> Option<Self::Item> {
39        let result = self.inner.next();
40        if result.is_some() {
41            self.remaining -= 1;
42        }
43        result
44    }
45
46    /// Required for implementing `ExactSizeIterator`.
47    fn size_hint(&self) -> (usize, Option<usize>) {
48        (self.remaining, Some(self.remaining))
49    }
50}
51
52impl<I, Item, F, const LEN: usize> ExactSizeIterator for ArrayFlatMap<I, Item, F, LEN>
53where
54    I: Iterator,
55    F: FnMut(I::Item) -> [Item; LEN],
56{
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use core::mem::size_of;
63
64    #[test]
65    fn test_array_flat_map() {
66        static TEST_CASES: &[(&[u16], fn(u16) -> [u8; 2], &[u8])] = &[
67            // Empty input
68            (&[], u16::to_be_bytes, &[]),
69            // Non-empty input.
70            (
71                &[0x0102, 0x0304, 0x0506],
72                u16::to_be_bytes,
73                &[1, 2, 3, 4, 5, 6],
74            ),
75            // Test with a different mapping function.
76            (
77                &[0x0102, 0x0304, 0x0506],
78                u16::to_le_bytes,
79                &[2, 1, 4, 3, 6, 5],
80            ),
81        ];
82        TEST_CASES.iter().copied().for_each(|(input, f, expected)| {
83            let mapped = ArrayFlatMap::new(input.iter().copied(), f).unwrap();
84            super::super::test::assert_iterator(mapped, expected);
85        });
86    }
87
88    // Does ArrayFlatMap::new() handle overflow correctly?
89    #[test]
90    fn test_array_flat_map_len_overflow() {
91        struct DownwardCounter {
92            remaining: usize,
93        }
94        impl Iterator for DownwardCounter {
95            type Item = usize;
96
97            fn next(&mut self) -> Option<Self::Item> {
98                if self.remaining > 0 {
99                    let result = self.remaining;
100                    self.remaining -= 1;
101                    Some(result)
102                } else {
103                    None
104                }
105            }
106
107            fn size_hint(&self) -> (usize, Option<usize>) {
108                (self.remaining, Some(self.remaining))
109            }
110        }
111        impl ExactSizeIterator for DownwardCounter {}
112
113        const MAX: usize = usize::MAX / size_of::<usize>();
114
115        static TEST_CASES: &[(usize, bool)] = &[(MAX, true), (MAX + 1, false)];
116        TEST_CASES.iter().copied().for_each(|(input_len, is_some)| {
117            let inner = DownwardCounter {
118                remaining: input_len,
119            };
120            let mapped = ArrayFlatMap::new(inner, usize::to_be_bytes);
121            assert_eq!(mapped.is_some(), is_some);
122            if let Some(mapped) = mapped {
123                assert_eq!(mapped.len(), input_len * size_of::<usize>());
124            }
125        });
126    }
127}
ring/polyfill/array_flat_map.rs

ring/polyfill/
array_flat_map.rs
