ring/polyfill/
array_flat_map.rs

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
63    #[test]
64    fn test_array_flat_map() {
65        static TEST_CASES: &[(&[u16], fn(u16) -> [u8; 2], &[u8])] = &[
66            // Empty input
67            (&[], u16::to_be_bytes, &[]),
68            // Non-empty input.
69            (
70                &[0x0102, 0x0304, 0x0506],
71                u16::to_be_bytes,
72                &[1, 2, 3, 4, 5, 6],
73            ),
74            // Test with a different mapping function.
75            (
76                &[0x0102, 0x0304, 0x0506],
77                u16::to_le_bytes,
78                &[2, 1, 4, 3, 6, 5],
79            ),
80        ];
81        TEST_CASES.iter().copied().for_each(|(input, f, expected)| {
82            let mapped = ArrayFlatMap::new(input.iter().copied(), f).unwrap();
83            super::super::test::assert_iterator(mapped, expected);
84        });
85    }
86
87    // Does ArrayFlatMap::new() handle overflow correctly?
88    #[test]
89    fn test_array_flat_map_len_overflow() {
90        struct DownwardCounter {
91            remaining: usize,
92        }
93        impl Iterator for DownwardCounter {
94            type Item = usize;
95
96            fn next(&mut self) -> Option<Self::Item> {
97                if self.remaining > 0 {
98                    let result = self.remaining;
99                    self.remaining -= 1;
100                    Some(result)
101                } else {
102                    None
103                }
104            }
105
106            fn size_hint(&self) -> (usize, Option<usize>) {
107                (self.remaining, Some(self.remaining))
108            }
109        }
110        impl ExactSizeIterator for DownwardCounter {}
111
112        const MAX: usize = usize::MAX / core::mem::size_of::<usize>();
113
114        static TEST_CASES: &[(usize, bool)] = &[(MAX, true), (MAX + 1, false)];
115        TEST_CASES.iter().copied().for_each(|(input_len, is_some)| {
116            let inner = DownwardCounter {
117                remaining: input_len,
118            };
119            let mapped = ArrayFlatMap::new(inner, usize::to_be_bytes);
120            assert_eq!(mapped.is_some(), is_some);
121            if let Some(mapped) = mapped {
122                assert_eq!(mapped.len(), input_len * core::mem::size_of::<usize>());
123            }
124        });
125    }
126}
ring/polyfill/array_flat_map.rs

ring/polyfill/
array_flat_map.rs
